From 3aa02eb8de191a226fdd768ff2e15099ddbbfa67 Mon Sep 17 00:00:00 2001 From: Israel Roldan Date: Fri, 19 Apr 2024 12:25:46 -0600 Subject: [PATCH] Added proposal to divide the main functionality of the Benchmark Base class. --- benchmarks/benchmark_base.py | 60 ++-------------- benchmarks/benchmark_template.py | 2 +- benchmarks/io_configuration_config_reader.py | 28 ++++---- benchmarks/plasma_nlte_excitation_x.py | 2 +- benchmarks/util/__init__.py | 0 benchmarks/util/base.py | 20 ++++++ benchmarks/util/nlte.py | 72 ++++++++++++++++++++ 7 files changed, 112 insertions(+), 72 deletions(-) create mode 100644 benchmarks/util/__init__.py create mode 100644 benchmarks/util/base.py create mode 100644 benchmarks/util/nlte.py diff --git a/benchmarks/benchmark_base.py b/benchmarks/benchmark_base.py index f8a584f5258..e84b54e5c6a 100644 --- a/benchmarks/benchmark_base.py +++ b/benchmarks/benchmark_base.py @@ -9,6 +9,7 @@ import pandas as pd from numba import njit +from benchmarks.util.nlte import Nlte from tardis.io.atom_data import AtomData from tardis.io.configuration import config_reader from tardis.io.configuration.config_reader import Configuration @@ -29,6 +30,9 @@ class BenchmarkBase: # 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__)) @@ -52,46 +56,6 @@ def tardis_config_verysimple(self): YAMLLoader, ) - @property - def tardis_config_verysimple_nlte(self): - filename = self.get_absolute_path("tardis/io/configuration/tests/data/tardis_configv1_nlte.yml") - return yaml_load_file( - filename, - YAMLLoader, - ) - - @property - def nlte_raw_model_root(self): - return SimulationState.from_config( - self.tardis_model_config_nlte_root, self.nlte_atom_data - ) - - @property - def nlte_raw_model_lu(self): - return SimulationState.from_config( - self.tardis_model_config_nlte_lu, self.nlte_atom_data - ) - - @property - def nlte_atom_data(self): - atomic_data = deepcopy(self.nlte_atomic_dataset) - return atomic_data - - @property - def nlte_atomic_dataset(self): - nlte_atomic_data = AtomData.from_hdf(self.nlte_atomic_data_fname) - return nlte_atomic_data - - @property - def nlte_atomic_data_fname(self): - atomic_data_fname = f"{self.tardis_ref_path}/nlte_atom_data/TestNLTE_He_Ti.h5" - atom_data_missing_str = f"{atomic_data_fname} atomic datafiles " f"does not seem to exist" - - if not Path(atomic_data_fname).exists(): - raise Exception(atom_data_missing_str) - - return atomic_data_fname - @property def tardis_ref_path(self): # TODO: This route is fixed but needs to get from the arguments given in the command line. @@ -118,22 +82,6 @@ def atomic_data_fname(self): return atomic_data_fname - @property - def tardis_model_config_nlte_root(self): - config = Configuration.from_yaml( - f"{self.example_configuration_dir}/tardis_configv1_nlte.yml" - ) - config.plasma.nlte_solver = "root" - return config - - @property - def tardis_model_config_nlte_lu(self): - config = Configuration.from_yaml( - f"{self.example_configuration_dir}/tardis_configv1_nlte.yml" - ) - config.plasma.nlte_solver = "lu" - return config - @property def example_configuration_dir(self): return self.get_absolute_path("tardis/io/configuration/tests/data") diff --git a/benchmarks/benchmark_template.py b/benchmarks/benchmark_template.py index 9210bd1ad4e..8aad9f4c40c 100644 --- a/benchmarks/benchmark_template.py +++ b/benchmarks/benchmark_template.py @@ -12,7 +12,7 @@ class BenchmarkXx(BenchmarkBase): """ def __init__(self): - pass + super().__init__() def time_template(self): pass diff --git a/benchmarks/io_configuration_config_reader.py b/benchmarks/io_configuration_config_reader.py index 9c42ea98e8f..aebf7abaf89 100644 --- a/benchmarks/io_configuration_config_reader.py +++ b/benchmarks/io_configuration_config_reader.py @@ -91,32 +91,32 @@ def time_plasma_section_config(self, key): ) def time_plasma_nlte_section_root_config(self): - self.tardis_config_verysimple_nlte["plasma"]["continuum_interaction"]["species"] = ["He I", ] - self.tardis_config_verysimple_nlte["plasma"]["nlte_ionization_species"] = ["H I"] - config = Configuration.from_config_dict(self.tardis_config_verysimple_nlte) - assemble_plasma(config, self.nlte_raw_model_root, self.nlte_atom_data) + self.nlte.tardis_config_verysimple_nlte["plasma"]["continuum_interaction"]["species"] = ["He I", ] + self.nlte.tardis_config_verysimple_nlte["plasma"]["nlte_ionization_species"] = ["H I"] + config = Configuration.from_config_dict(self.nlte.tardis_config_verysimple_nlte) + assemble_plasma(config, self.nlte.nlte_raw_model_root, self.nlte.nlte_atom_data) def time_plasma_nlte_section_lu_config(self): - self.tardis_config_verysimple_nlte["plasma"]["continuum_interaction"][ + self.nlte.tardis_config_verysimple_nlte["plasma"]["continuum_interaction"][ "species" ] = [ "He I", ] - self.tardis_config_verysimple_nlte["plasma"]["nlte_ionization_species"] = ["H I"] - self.tardis_config_verysimple_nlte["plasma"]["nlte_solver"] = "lu" - config = Configuration.from_config_dict(self.tardis_config_verysimple_nlte) - assemble_plasma(config, self.nlte_raw_model_lu, self.nlte_atom_data) + self.nlte.tardis_config_verysimple_nlte["plasma"]["nlte_ionization_species"] = ["H I"] + self.nlte.tardis_config_verysimple_nlte["plasma"]["nlte_solver"] = "lu" + config = Configuration.from_config_dict(self.nlte.tardis_config_verysimple_nlte) + assemble_plasma(config, self.nlte.nlte_raw_model_lu, self.nlte.nlte_atom_data) def time_plasma_nlte_root_exc_section_config(self): - self.tardis_config_verysimple_nlte["plasma"]["continuum_interaction"][ + self.nlte.tardis_config_verysimple_nlte["plasma"]["continuum_interaction"][ "species" ] = [ "He I", ] - self.tardis_config_verysimple_nlte["plasma"]["nlte_excitation_species"] = ["H I"] - self.tardis_config_verysimple_nlte["plasma"]["nlte_solver"] = "root" - config = Configuration.from_config_dict(self.tardis_config_verysimple_nlte) - assemble_plasma(config, self.nlte_raw_model_root, self.nlte_atom_data) + self.nlte.tardis_config_verysimple_nlte["plasma"]["nlte_excitation_species"] = ["H I"] + self.nlte.tardis_config_verysimple_nlte["plasma"]["nlte_solver"] = "root" + config = Configuration.from_config_dict(self.nlte.tardis_config_verysimple_nlte) + assemble_plasma(config, self.nlte.nlte_raw_model_root, self.nlte.nlte_atom_data) def time_spectrum_section_config(self): self.tardis_config_verysimple["spectrum"]["start"] = Quantity("2500 angstrom") diff --git a/benchmarks/plasma_nlte_excitation_x.py b/benchmarks/plasma_nlte_excitation_x.py index a6f1f1a2725..963620d5224 100644 --- a/benchmarks/plasma_nlte_excitation_x.py +++ b/benchmarks/plasma_nlte_excitation_x.py @@ -22,7 +22,7 @@ class BenchmarkPlasmaNlteExcitation(BenchmarkBase): """ def time_prepare_bound_bound_rate_matrix(self): - nlte_atomic_dataset = self.nlte_atomic_dataset + nlte_atomic_dataset = self.nlte.nlte_atomic_dataset # TODO: Needs to work in the class RegressionData in BenchmarkBase custom_request = self.CustomPyTestRequest( tardis_regression_data_path="/app/tardis-regression-data", diff --git a/benchmarks/util/__init__.py b/benchmarks/util/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/benchmarks/util/base.py b/benchmarks/util/base.py new file mode 100644 index 00000000000..e1007ebbf2e --- /dev/null +++ b/benchmarks/util/base.py @@ -0,0 +1,20 @@ +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 new file mode 100644 index 00000000000..4d91a7c78bc --- /dev/null +++ b/benchmarks/util/nlte.py @@ -0,0 +1,72 @@ +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