From ab98e618580e8e1efe39715d231cbba499d09cf9 Mon Sep 17 00:00:00 2001 From: Israel Roldan Date: Fri, 19 Apr 2024 04:48:32 -0600 Subject: [PATCH] Fixed the usage of the regression data class. --- benchmarks/benchmark_base.py | 188 +++++------------ benchmarks/plasma_nlte_excitation_x.py | 271 +++++++++++++------------ 2 files changed, 187 insertions(+), 272 deletions(-) diff --git a/benchmarks/benchmark_base.py b/benchmarks/benchmark_base.py index c4498396b74..f8a584f5258 100644 --- a/benchmarks/benchmark_base.py +++ b/benchmarks/benchmark_base.py @@ -21,6 +21,7 @@ ) from tardis.simulation import Simulation from tardis.tests.fixtures.atom_data import DEFAULT_ATOM_DATA_UUID +from tardis.tests.fixtures.regression_data import RegressionData class BenchmarkBase: @@ -198,157 +199,64 @@ def config_verysimple(self): f"{self.example_configuration_dir}/tardis_configv1_verysimple.yml" ) - class RegressionData: - def __init__(self) -> None: - # TODO: This route is fixed but needs to get from the arguments given in the command line. - # /app/tardis-regression-data - self.regression_data_path = "/app/tardis-regression-data" - # TODO: Parameter `--generate-reference` is set in the command line. - # It is fixed but needs to get from the arguments. - self.enable_generate_reference = False - self.fname = f"{self.fname_prefix}.UNKNOWN_FORMAT" + 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 module_name(self): - return self.__name__ + def config(self): + class SubClass: + @staticmethod + def getoption(option): + if option == '--tardis-regression-data': + return self.tardis_regression_data_path + return None - @property - def test_name(self): - return self.name + return SubClass() @property - def fname_prefix(self): - double_under = re.compile(r"[:\[\]{}]") - no_space = re.compile(r'[,"\']') # quotes and commas + def node(self): + class SubClass: + def __init__(self, parent): + self.parent = parent - name = double_under.sub("__", self.test_name) - name = no_space.sub("", name) - return name + @property + def name(self): + return self.parent.node_name - @property - def relative_regression_data_dir(self): - relative_data_dir = Path(self.module_name.replace(".", "/")) - if self.cls is not None: - relative_data_dir /= HDFWriterMixin.convert_to_snake_case( - self.cls.__name__ - ) - return relative_data_dir + @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 absolute_regression_data_dir(self): - return f"{self.regression_data_path}/{self.relative_regression_data_dir}" + def cls(self): + return None @property - def fpath(self): - return f"{self.absolute_regression_data_dir}/{self.fname}" - - def sync_dataframe(self, data, key="data"): - """ - Synchronizes the dataframe with the regression data. - - Parameters - ---------- - data : DataFrame - The dataframe to be synchronized. - key : str, optional - The key to use for storing the dataframe in the regression data file. Defaults to "data". - - Returns - ------- - DataFrame or None - The synchronized dataframe if `enable_generate_reference` is `False`, otherwise `None`. - """ - self.fname = f"{self.fname_prefix}.h5" - if self.enable_generate_reference: - Path(self.fpath).parent.mkdir(parents=True, exist_ok=True) - data.to_hdf( - self.fpath, - key=key, - ) - raise Exception("Skipping test to generate reference data") - else: - return pd.read_hdf(self.fpath, key=key) - - def sync_ndarray(self, data): - """ - Synchronizes the ndarray with the regression data. - - Parameters - ---------- - data : ndarray - The ndarray to be synchronized. - - Returns - ------- - ndarray or None - The synchronized ndarray if `enable_generate_reference` is `False`, otherwise `None`. - """ - self.fname = f"{self.fname_prefix}.npy" - if self.enable_generate_reference: - Path(self.fpath).parent.mkdir(parents=True, exist_ok=True) - Path(self.fpath).parent.mkdir(parents=True, exist_ok=True) - np.save(self.fpath, data) - raise Exception("Skipping test to generate reference data") - else: - return np.load(self.fpath) - - def sync_str(self, data): - """ - Synchronizes the string with the regression data. - - Parameters - ---------- - data : str - The string to be synchronized. - - Returns - ------- - str or None - The synchronized string if `enable_generate_reference` is `False`, otherwise `None`. - """ - self.fname = f"{self.fname_prefix}.txt" - if self.enable_generate_reference: - Path(self.fpath).parent.mkdir(parents=True, exist_ok=True) - with Path(self.fpath).open("w") as fh: - fh.write(data) - raise Exception( - f"Skipping test to generate regression_data {self.fpath} data" - ) - else: - with Path(self.fpath).open("r") as fh: - return fh.read() - - def sync_hdf_store(self, tardis_module, update_fname=True): - """ - Synchronizes the HDF store with the regression data. - - Parameters - ---------- - tardis_module : object - The module to be synchronized. - update_fname : bool, optional - Whether to update the file name. Defaults to True. - - Returns - ------- - HDFStore or None - The synchronized HDF store if `enable_generate_reference` is `False`, otherwise `None`. - """ - if update_fname: - self.fname = f"{self.fname_prefix}.h5" - if self.enable_generate_reference: - Path(self.fpath).parent.mkdir(parents=True, exist_ok=True) - with pd.HDFStore(self.fpath, mode="w") as store: - tardis_module.to_hdf(store, overwrite=True) - raise Exception( - f"Skipping test to generate regression_data {self.fpath} data" - ) - else: - return pd.HDFStore(self.fpath, mode="r") - - @property - def regression_data(self): - return self.RegressionData() + def relative_regression_data_dir(self): + return self.regression_data_dir + + @staticmethod + def regression_data(request: CustomPyTestRequest): + return RegressionData(request) @property def packet(self): diff --git a/benchmarks/plasma_nlte_excitation_x.py b/benchmarks/plasma_nlte_excitation_x.py index e5f963dded5..a6f1f1a2725 100644 --- a/benchmarks/plasma_nlte_excitation_x.py +++ b/benchmarks/plasma_nlte_excitation_x.py @@ -1,135 +1,142 @@ """ Basic TARDIS Benchmark. """ -# import numpy as np -# import numpy.testing as npt -# import pandas as pd -# from asv_runner.benchmarks.mark import parameterize, skip_benchmark -# -# from benchmarks.benchmark_base import BenchmarkBase -# from tardis.plasma.properties.nlte_excitation_data import NLTEExcitationData -# from tardis.plasma.properties.nlte_rate_equation_solver import ( -# prepare_r_uls_r_lus, -# prepare_bound_bound_rate_matrix, -# create_coll_exc_deexc_matrix, -# ) -# -# -# # @skip_benchmark -# class BenchmarkPlasmaNlteExcitation(BenchmarkBase): -# """ -# Class to benchmark the NLTE excitation function. -# """ -# -# def __init__(self): -# pass -# -# def time_prepare_bound_bound_rate_matrix(self): -# nlte_atomic_dataset = self.nlte_atomic_dataset -# # TODO: Needs to work in the class RegressionData in BenchmarkBase -# regression_data = self.regression_data -# simple_excitation_species = [(1, 0)] -# copy_atomic_dataset = nlte_atomic_dataset -# copy_atomic_dataset.levels = nlte_atomic_dataset.levels[ -# nlte_atomic_dataset.levels.index.get_level_values("level_number") < 5 -# ] -# lines_filtered = nlte_atomic_dataset.lines[ -# nlte_atomic_dataset.lines.index.get_level_values("level_number_lower") -# < 5 -# ] -# copy_atomic_dataset.lines = lines_filtered[ -# lines_filtered.index.get_level_values("level_number_upper") < 5 -# ] -# simple_nlte_data = NLTEExcitationData( -# copy_atomic_dataset.lines, simple_excitation_species -# ) -# simple_number_of_shells = 1 -# simple_j_blues_matrix = np.linspace( -# 0.1, 1.0, copy_atomic_dataset.lines.index.size -# ) -# simple_j_blues = pd.DataFrame( -# simple_j_blues_matrix, -# index=copy_atomic_dataset.lines.index, -# columns=["0"], -# ) -# simple_number_of_levels = copy_atomic_dataset.levels.energy.loc[ -# simple_excitation_species[0] -# ].count() -# ( -# lines_index, -# r_ul_index, -# r_ul_matrix, -# r_lu_index, -# r_lu_matrix, -# ) = prepare_r_uls_r_lus( -# simple_number_of_levels, -# simple_number_of_shells, -# simple_j_blues, -# simple_excitation_species[0], -# simple_nlte_data, -# ) -# simple_beta_sobolev_matrix = np.linspace( -# 2.5, 4.7, copy_atomic_dataset.lines.index.size -# ) -# simple_beta_sobolev = pd.DataFrame( -# simple_beta_sobolev_matrix, -# index=copy_atomic_dataset.lines.index, -# columns=["0"], -# ) -# actual_rate_matrix = prepare_bound_bound_rate_matrix( -# simple_number_of_levels, -# lines_index, -# r_ul_index, -# r_ul_matrix, -# r_lu_index, -# r_lu_matrix, -# simple_beta_sobolev, -# ) -# # if this test fails the first thing to check is if the reshape in the -# # methods made a view or a copy. If it's a copy rewrite the function. -# # TODO: allow rtol=1e-6 -# expected_rate_matrix = regression_data.sync_ndarray(actual_rate_matrix) -# npt.assert_allclose(actual_rate_matrix, expected_rate_matrix, rtol=1e-6) -# -# @parameterize( -# [ -# "coll_exc_coeff_values", -# "coll_deexc_coeff_values", -# "number_of_levels", -# ], -# [ -# ( -# [1, -2, 3], -# [4, 9, 10], -# 3, -# ), -# ( -# [0.21, 0.045, 0.1234], -# [0.7865, 0.987, 0.00123], -# 3, -# ), -# ], -# ) -# def time_coll_exc_deexc_matrix( -# coll_exc_coeff_values, -# coll_deexc_coeff_values, -# number_of_levels, -# regression_data, -# ): -# """ -# Checks the NLTERateEquationSolver.create_coll_exc_deexc_matrix for simple values of species with 3 levels. -# NOTE: Values used for testing are not physical. -# """ -# index = pd.MultiIndex.from_tuples( -# [(0, 1), (0, 2), (1, 2)], -# names=["level_number_lower", "level_number_upper"], -# ) -# exc_coeff = pd.Series(coll_exc_coeff_values, index=index) -# deexc_coeff = pd.Series(coll_deexc_coeff_values, index=index) -# obtained_coeff_matrix = create_coll_exc_deexc_matrix( -# exc_coeff, deexc_coeff, number_of_levels -# ) -# expected_obtained_coeff_matrix = regression_data.sync_ndarray( -# obtained_coeff_matrix -# ) -# npt.assert_allclose(expected_obtained_coeff_matrix, obtained_coeff_matrix) +import numpy as np +import numpy.testing as npt +import pandas as pd +from asv_runner.benchmarks.mark import parameterize, skip_benchmark + +from benchmarks.benchmark_base import BenchmarkBase +from tardis.plasma.properties.nlte_excitation_data import NLTEExcitationData +from tardis.plasma.properties.nlte_rate_equation_solver import ( + prepare_r_uls_r_lus, + prepare_bound_bound_rate_matrix, + create_coll_exc_deexc_matrix, +) + + +# @skip_benchmark +class BenchmarkPlasmaNlteExcitation(BenchmarkBase): + """ + Class to benchmark the NLTE excitation function. + """ + + def time_prepare_bound_bound_rate_matrix(self): + nlte_atomic_dataset = self.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", + node_name="test_prepare_bound_bound_rate_matrix", + node_module_name="tardis.plasma.tests.test_nlte_excitation", + regression_data_dir="tardis/plasma/tests/test_nlte_excitation", + ) + regression_data = self.regression_data(custom_request) + + simple_excitation_species = [(1, 0)] + copy_atomic_dataset = nlte_atomic_dataset + copy_atomic_dataset.levels = nlte_atomic_dataset.levels[ + nlte_atomic_dataset.levels.index.get_level_values("level_number") < 5 + ] + lines_filtered = nlte_atomic_dataset.lines[ + nlte_atomic_dataset.lines.index.get_level_values("level_number_lower") + < 5 + ] + copy_atomic_dataset.lines = lines_filtered[ + lines_filtered.index.get_level_values("level_number_upper") < 5 + ] + simple_nlte_data = NLTEExcitationData( + copy_atomic_dataset.lines, simple_excitation_species + ) + simple_number_of_shells = 1 + simple_j_blues_matrix = np.linspace( + 0.1, 1.0, copy_atomic_dataset.lines.index.size + ) + simple_j_blues = pd.DataFrame( + simple_j_blues_matrix, + index=copy_atomic_dataset.lines.index, + columns=["0"], + ) + simple_number_of_levels = copy_atomic_dataset.levels.energy.loc[ + simple_excitation_species[0] + ].count() + ( + lines_index, + r_ul_index, + r_ul_matrix, + r_lu_index, + r_lu_matrix, + ) = prepare_r_uls_r_lus( + simple_number_of_levels, + simple_number_of_shells, + simple_j_blues, + simple_excitation_species[0], + simple_nlte_data, + ) + simple_beta_sobolev_matrix = np.linspace( + 2.5, 4.7, copy_atomic_dataset.lines.index.size + ) + simple_beta_sobolev = pd.DataFrame( + simple_beta_sobolev_matrix, + index=copy_atomic_dataset.lines.index, + columns=["0"], + ) + actual_rate_matrix = prepare_bound_bound_rate_matrix( + simple_number_of_levels, + lines_index, + r_ul_index, + r_ul_matrix, + r_lu_index, + r_lu_matrix, + simple_beta_sobolev, + ) + # if this test fails the first thing to check is if the reshape in the + # methods made a view or a copy. If it's a copy rewrite the function. + # TODO: allow rtol=1e-6 + expected_rate_matrix = regression_data.sync_ndarray(actual_rate_matrix) + npt.assert_allclose(actual_rate_matrix, expected_rate_matrix, rtol=1e-6) + + @parameterize({'Matrix': [ + { + "coll_exc_coeff_values": [1, -2, 3], + "coll_deexc_coeff_values": [4, 9, 10], + "number_of_levels": 3, + }, + { + "coll_exc_coeff_values": [0.21, 0.045, 0.1234], + "coll_deexc_coeff_values": [0.7865, 0.987, 0.00123], + "number_of_levels": 3, + }, + ]}) + def time_coll_exc_deexc_matrix(self, matrix): + """ + Checks the NLTERateEquationSolver.create_coll_exc_deexc_matrix for simple values of species with 3 levels. + NOTE: Values used for testing are not physical. + """ + coll_exc_coeff_values = matrix["coll_exc_coeff_values"] + coll_deexc_coeff_values = matrix["coll_deexc_coeff_values"] + number_of_levels = matrix["number_of_levels"] + index = pd.MultiIndex.from_tuples( + [(0, 1), (0, 2), (1, 2)], + names=["level_number_lower", "level_number_upper"], + ) + exc_coeff = pd.Series(coll_exc_coeff_values, index=index) + deexc_coeff = pd.Series(coll_deexc_coeff_values, index=index) + obtained_coeff_matrix = create_coll_exc_deexc_matrix( + exc_coeff, deexc_coeff, number_of_levels + ) + # TODO: Needs to work in the class RegressionData in BenchmarkBase + test_parameter_id = 0 + if coll_exc_coeff_values != [1, -2, 3]: + test_parameter_id = 1 + custom_request = self.CustomPyTestRequest( + "/app/tardis-regression-data", + f"test_coll_exc_deexc_matrix__coll_exc_coeff_values{test_parameter_id}-coll_deexc_coeff_values{test_parameter_id}-3__", + "tardis.plasma.tests.test_nlte_excitation", + "tardis/plasma/tests/test_nlte_excitation", + ) + regression_data = self.regression_data(custom_request) + + expected_obtained_coeff_matrix = regression_data.sync_ndarray( + obtained_coeff_matrix + ) + npt.assert_allclose(expected_obtained_coeff_matrix, obtained_coeff_matrix)