Added proposal to divide the main functionality of the Benchmark Base…

… class.

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")

benchmarks/benchmark_template.py

@@ -12,7 +12,7 @@ class BenchmarkXx(BenchmarkBase):
     """
 
     def __init__(self):
-        pass
+        super().__init__()
 
     def time_template(self):
         pass

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")

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",

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")

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