From 59c6e3f85bb2d811571f90d6c9af633992979d41 Mon Sep 17 00:00:00 2001 From: rohchav Date: Wed, 19 Feb 2025 12:55:29 -0600 Subject: [PATCH 1/4] refactoring scattering angle calculation --- test_1d.py | 61 +++++++++++++++++++++++ tests/test_forward/test_1d.py | 25 +++++----- tests/test_forward/test_angular_1d.py | 39 ++++++++------- tests/test_inverse/test_1d_random.py | 31 +++++++----- tsadar/core/thomson_diagnostic.py | 53 ++++++++++++++++++-- tsadar/inverse/loss_function.py | 2 +- tsadar/utils/data_handling/calibration.py | 2 +- 7 files changed, 163 insertions(+), 50 deletions(-) create mode 100644 test_1d.py diff --git a/test_1d.py b/test_1d.py new file mode 100644 index 00000000..e982df64 --- /dev/null +++ b/test_1d.py @@ -0,0 +1,61 @@ +import time, pytest +import multiprocessing as mp +import yaml +import mlflow +from flatten_dict import flatten, unflatten +from numpy.testing import assert_allclose +from jax import config + +config.update("jax_enable_x64", True) +# config.update("jax_disable_jit", True) +# config.update("jax_check_tracer_leaks", True) + +from tsadar.inverse import fitter +from tsadar.utils import misc + + +@pytest.mark.parametrize("nn", [False]) +def test_data(nn): + # Test #3: Data test, compare fit to a preknown fit result + # currently just runs one line of shot 101675 for the electron, should be expanded in the future + + with open("tests/configs/time_test_defaults.yaml", "r") as fi: + defaults = yaml.safe_load(fi) + + with open("tests/configs/time_test_inputs.yaml", "r") as fi: + inputs = yaml.safe_load(fi) + + defaults = flatten(defaults) + defaults.update(flatten(inputs)) + config = unflatten(defaults) + + config["nn"]["use"] = nn + # config["parameters"]["Te"]["val"] = 0.5 + # config["parameters"]["ne"]["val"] = 0.2 # 0.25 + # config["parameters"]["m"]["val"] = 3.0 # 2.2 + + mlflow.set_experiment(config["mlflow"]["experiment"]) + + with mlflow.start_run() as run: + misc.log_mlflow(config) + config["num_cores"] = int(mp.cpu_count()) + + t0 = time.time() + fit_results, loss = fitter.fit(config=config) + metrics_dict = {"total_time": time.time() - t0, "num_cores": int(mp.cpu_count())} + mlflow.log_metrics(metrics=metrics_dict) + mlflow.set_tag("status", "completed") + print(fit_results) + + # These were changed 5/6/24 to reflect new good fit values, unclear why changes were required probably a change + # to the calibration + assert_allclose(fit_results["amp1_general"][0], 0.734, rtol=1e-1) # 0.9257 + assert_allclose(fit_results["amp2_general"][0], 0.519, rtol=1e-1) # 0.6727 + assert_allclose(fit_results["lam_general"][0], 524.016, rtol=5e-3) # 524.2455 + assert_allclose(fit_results["Te_electron"][0], 0.5994, rtol=1e-1) # 0.67585 + assert_allclose(fit_results["ne_electron"][0], 0.2256, rtol=5e-2) # 0.21792 + assert_allclose(fit_results["m_electron"][0], 2.987, rtol=15e-2) # 3.3673 + + +if __name__ == "__main__": + test_data(False) diff --git a/tests/test_forward/test_1d.py b/tests/test_forward/test_1d.py index 72faef65..c69fae32 100644 --- a/tests/test_forward/test_1d.py +++ b/tests/test_forward/test_1d.py @@ -11,7 +11,7 @@ from tsadar.utils import misc from tsadar.core.thomson_diagnostic import ThomsonScatteringDiagnostic from tsadar.core.modules import ThomsonParams -from tsadar.utils.data_handling.calibration import get_scattering_angles +# from tsadar.utils.data_handling.calibration import get_scattering_angles def test_1d_forward_pass(): @@ -40,16 +40,16 @@ def test_1d_forward_pass(): config = unflatten(defaults) # get scattering angles and weights - config["other"]["lamrangE"] = [ - config["data"]["fit_rng"]["forward_epw_start"], - config["data"]["fit_rng"]["forward_epw_end"], - ] - config["other"]["lamrangI"] = [ - config["data"]["fit_rng"]["forward_iaw_start"], - config["data"]["fit_rng"]["forward_iaw_end"], - ] - config["other"]["npts"] = int(config["other"]["CCDsize"][1] * config["other"]["points_per_pixel"]) - sas = get_scattering_angles(config) + # config["other"]["lamrangE"] = [ + # config["data"]["fit_rng"]["forward_epw_start"], + # config["data"]["fit_rng"]["forward_epw_end"], + # ] + # config["other"]["lamrangI"] = [ + # config["data"]["fit_rng"]["forward_iaw_start"], + # config["data"]["fit_rng"]["forward_iaw_end"], + # ] + # config["other"]["npts"] = int(config["other"]["CCDsize"][1] * config["other"]["points_per_pixel"]) + # sas = get_scattering_angles(config) dummy_batch = { "i_data": np.array([1]), @@ -60,7 +60,8 @@ def test_1d_forward_pass(): "i_amps": np.array([1]), } - ts_diag = ThomsonScatteringDiagnostic(config, scattering_angles=sas) + ts_diag = ThomsonScatteringDiagnostic(config, scattering_angles=None, angular=False) + config = ts_diag.get_cfg() ts_params = ThomsonParams(config["parameters"], num_params=1, batch=True, activate=True) ThryE, ThryI, lamAxisE, lamAxisI = ts_diag(ts_params, dummy_batch) diff --git a/tests/test_forward/test_angular_1d.py b/tests/test_forward/test_angular_1d.py index a9ab9370..8b6ba66a 100644 --- a/tests/test_forward/test_angular_1d.py +++ b/tests/test_forward/test_angular_1d.py @@ -11,7 +11,7 @@ from tsadar.core.thomson_diagnostic import ThomsonScatteringDiagnostic from tsadar.core.modules import ThomsonParams -from tsadar.utils.data_handling.calibration import get_scattering_angles, get_calibrations +# from tsadar.utils.data_handling.calibration import get_scattering_angles, get_calibrations def test_arts1d_forward_pass(): @@ -40,23 +40,23 @@ def test_arts1d_forward_pass(): config = unflatten(defaults) # get scattering angles and weights - config["other"]["lamrangE"] = [ - config["data"]["fit_rng"]["forward_epw_start"], - config["data"]["fit_rng"]["forward_epw_end"], - ] - config["other"]["lamrangI"] = [ - config["data"]["fit_rng"]["forward_iaw_start"], - config["data"]["fit_rng"]["forward_iaw_end"], - ] - config["other"]["npts"] = int(config["other"]["CCDsize"][1] * config["other"]["points_per_pixel"]) - sas = get_scattering_angles(config) - - [axisxE, _, _, _, _, _] = get_calibrations( - 104000, config["other"]["extraoptions"]["spectype"], 0.0, config["other"]["CCDsize"] - ) # shot number hardcoded to get calibration - config["other"]["extraoptions"]["spectype"] = "angular_full" - - sas["angAxis"] = axisxE + # config["other"]["lamrangE"] = [ + # config["data"]["fit_rng"]["forward_epw_start"], + # config["data"]["fit_rng"]["forward_epw_end"], + # ] + # config["other"]["lamrangI"] = [ + # config["data"]["fit_rng"]["forward_iaw_start"], + # config["data"]["fit_rng"]["forward_iaw_end"], + # ] + # config["other"]["npts"] = int(config["other"]["CCDsize"][1] * config["other"]["points_per_pixel"]) + # sas = get_scattering_angles(config) + + # [axisxE, _, _, _, _, _] = get_calibrations( + # 104000, config["other"]["extraoptions"]["spectype"], 0.0, config["other"]["CCDsize"] + # ) # shot number hardcoded to get calibration + # config["other"]["extraoptions"]["spectype"] = "angular_full" + + # sas["angAxis"] = axisxE dummy_batch = { "i_data": np.ones((config["other"]["CCDsize"][0], config["other"]["CCDsize"][1])), @@ -67,7 +67,8 @@ def test_arts1d_forward_pass(): "i_amps": np.array([1]), } - ts_diag = ThomsonScatteringDiagnostic(config, scattering_angles=sas) + ts_diag = ThomsonScatteringDiagnostic(config, scattering_angles=None, angular=True) + config = ts_diag.get_cfg() ts_params = ThomsonParams(config["parameters"], num_params=1, batch=False, activate=True) ThryE, ThryI, lamAxisE, lamAxisI = ts_diag(ts_params, dummy_batch) # np.save("tests/test_forward/ThryE-arts1d.npy", ThryE) diff --git a/tests/test_inverse/test_1d_random.py b/tests/test_inverse/test_1d_random.py index 2675f9a1..ea15b8a7 100644 --- a/tests/test_inverse/test_1d_random.py +++ b/tests/test_inverse/test_1d_random.py @@ -14,7 +14,7 @@ from tsadar.utils import misc from tsadar.core.thomson_diagnostic import ThomsonScatteringDiagnostic from tsadar.core.modules import ThomsonParams, get_filter_spec -from tsadar.utils.data_handling.calibration import get_scattering_angles +# from tsadar.utils.data_handling.calibration import get_scattering_angles def _perturb_params_(rng, params): @@ -80,16 +80,16 @@ def test_1d_inverse(): config = unflatten(defaults) # get scattering angles and weights - config["other"]["lamrangE"] = [ - config["data"]["fit_rng"]["forward_epw_start"], - config["data"]["fit_rng"]["forward_epw_end"], - ] - config["other"]["lamrangI"] = [ - config["data"]["fit_rng"]["forward_iaw_start"], - config["data"]["fit_rng"]["forward_iaw_end"], - ] - config["other"]["npts"] = int(config["other"]["CCDsize"][1] * config["other"]["points_per_pixel"]) - sas = get_scattering_angles(config) + # config["other"]["lamrangE"] = [ + # config["data"]["fit_rng"]["forward_epw_start"], + # config["data"]["fit_rng"]["forward_epw_end"], + # ] + # config["other"]["lamrangI"] = [ + # config["data"]["fit_rng"]["forward_iaw_start"], + # config["data"]["fit_rng"]["forward_iaw_end"], + # ] + # config["other"]["npts"] = int(config["other"]["CCDsize"][1] * config["other"]["points_per_pixel"]) + # sas = get_scattering_angles(config) dummy_batch = { "i_data": np.array([1]), @@ -100,7 +100,8 @@ def test_1d_inverse(): "i_amps": np.array([1]), } rng = np.random.default_rng() - ts_diag = ThomsonScatteringDiagnostic(config, scattering_angles=sas) + ts_diag = ThomsonScatteringDiagnostic(config, scattering_angles=None, angular=False) + config = ts_diag.get_cfg() config["parameters"] = _perturb_params_(rng, config["parameters"]) misc.log_mlflow(config) ts_params_gt = ThomsonParams(config["parameters"], num_params=1, batch=True, activate=True) @@ -109,7 +110,7 @@ def test_1d_inverse(): loss = 1 while np.nan_to_num(loss, nan=1) > 1e-3: - ts_diag = ThomsonScatteringDiagnostic(config, scattering_angles=sas) + ts_diag = ThomsonScatteringDiagnostic(config, scattering_angles=None, angular=False) config["parameters"] = _perturb_params_(rng, config["parameters"]) ts_params_fit = ThomsonParams(config["parameters"], num_params=1, batch=True, activate=True) diff_params, static_params = eqx.partition( @@ -162,6 +163,10 @@ def scipy_vg_fn(diff_params_flat): ax.set_ylabel("Intensity (arb. units)") ax.set_title("Electron Spectrum") fig.savefig(os.path.join(td, "ThryE.png"), bbox_inches="tight") + + # Save ThryE and ground_truth as text files + np.savetxt(os.path.join(td, "ThryE.txt"), ThryE) + np.savetxt(os.path.join(td, "ground_truth_ThryE.txt"), ground_truth["ThryE"]) mlflow.log_artifacts(td) # np.testing.assert_allclose(ThryE, ground_truth["ThryE"], atol=0, rtol=0.2) diff --git a/tsadar/core/thomson_diagnostic.py b/tsadar/core/thomson_diagnostic.py index 83b9fcb0..9a3ed247 100644 --- a/tsadar/core/thomson_diagnostic.py +++ b/tsadar/core/thomson_diagnostic.py @@ -1,11 +1,13 @@ from jax import numpy as jnp, vmap +from tsadar.utils.data_handling.calibration import get_scattering_angles, get_calibrations from .modules import ThomsonParams from .physics import irf from .physics.generate_spectra import FitModel + class ThomsonScatteringDiagnostic: """ The SpectrumCalculator class wraps the FitModel class adding instrumental effects to the calculated spectrum so it @@ -20,11 +22,13 @@ class ThomsonScatteringDiagnostic: weights of each of the scattering angles in the final spectrum """ - def __init__(self, cfg, scattering_angles): + def __init__(self, cfg, angular=False): + super().__init__() - self.cfg = cfg - self.scattering_angles = scattering_angles - self.model = FitModel(cfg, scattering_angles) + + self.cfg, self.scattering_angles = self.initialize_scattering_angles(cfg, angular) + + self.model = FitModel(cfg, self.scattering_angles) if ( "temporal" in cfg["other"]["extraoptions"]["spectype"] @@ -125,3 +129,44 @@ def __call__(self, ts_params: ThomsonParams, batch): ThryI = ThryI + batch["noise_i"] return ThryE, ThryI, lamAxisE, lamAxisI + + def initialize_scattering_angles(self, config, angular): + """ + Initializes scattering angles and weights. + + Args: + config: Configuration dictionary + + Returns: + Updated configuration dictionary with scattering angles and weights + """ + config["other"]["lamrangE"] = [ + config["data"]["fit_rng"]["forward_epw_start"], + config["data"]["fit_rng"]["forward_epw_end"], + ] + config["other"]["lamrangI"] = [ + config["data"]["fit_rng"]["forward_iaw_start"], + config["data"]["fit_rng"]["forward_iaw_end"], + ] + config["other"]["npts"] = int(config["other"]["CCDsize"][1] * config["other"]["points_per_pixel"]) + scattering_angles = get_scattering_angles(config) + + if angular: + [axisxE, _, _, _, _, _] = get_calibrations( + 104000, config["other"]["extraoptions"]["spectype"], 0.0, config["other"]["CCDsize"] + ) # shot number hardcoded to get calibration + config["other"]["extraoptions"]["spectype"] = "angular_full" + + scattering_angles["angAxis"] = axisxE + + return config, scattering_angles + + def get_cfg(self): + """ + Getter method for the cfg attribute + + Returns: + The configuration dictionary + """ + + return self.cfg \ No newline at end of file diff --git a/tsadar/inverse/loss_function.py b/tsadar/inverse/loss_function.py index 50b962a8..46ece644 100644 --- a/tsadar/inverse/loss_function.py +++ b/tsadar/inverse/loss_function.py @@ -45,7 +45,7 @@ def __init__(self, cfg: Dict, scattering_angles, dummy_batch): ############ - self.ts_diag = ThomsonScatteringDiagnostic(cfg, scattering_angles=scattering_angles) + self.ts_diag = ThomsonScatteringDiagnostic(cfg, angular=False) self._loss_ = filter_jit(self.__loss__) self._vg_func_ = filter_jit(filter_value_and_grad(self.__loss__, has_aux=True)) diff --git a/tsadar/utils/data_handling/calibration.py b/tsadar/utils/data_handling/calibration.py index b46f502a..1ea8ee67 100644 --- a/tsadar/utils/data_handling/calibration.py +++ b/tsadar/utils/data_handling/calibration.py @@ -189,7 +189,7 @@ def sa_lookup(beam): ), ) else: - raise NotImplmentedError("Other probe geometrries are not yet supported") + raise NotImplementedError("Other probe geometrries are not yet supported") return sa From 88b77011fbe2abc24319e9e3a4c3a411b6fe533b Mon Sep 17 00:00:00 2001 From: rohchav Date: Sat, 22 Feb 2025 13:39:01 -0600 Subject: [PATCH 2/4] fixed refactoring --- .gitignore | 1 + tests/test_forward/test_1d.py | 2 +- tests/test_forward/test_angular_1d.py | 2 +- tests/test_forward/test_angular_2d.py | 2 +- tests/test_inverse/test_1d_random.py | 4 +- tsadar/core/thomson_diagnostic.py | 62 ++++++++++++++++++--------- 6 files changed, 47 insertions(+), 26 deletions(-) diff --git a/.gitignore b/.gitignore index 39f40812..f36798b7 100644 --- a/.gitignore +++ b/.gitignore @@ -237,3 +237,4 @@ fabric.properties /mlruns .vscode/launch.json +myenv/ diff --git a/tests/test_forward/test_1d.py b/tests/test_forward/test_1d.py index c69fae32..c4d2198d 100644 --- a/tests/test_forward/test_1d.py +++ b/tests/test_forward/test_1d.py @@ -60,7 +60,7 @@ def test_1d_forward_pass(): "i_amps": np.array([1]), } - ts_diag = ThomsonScatteringDiagnostic(config, scattering_angles=None, angular=False) + ts_diag = ThomsonScatteringDiagnostic(config, angular=False, cumulative=False) config = ts_diag.get_cfg() ts_params = ThomsonParams(config["parameters"], num_params=1, batch=True, activate=True) ThryE, ThryI, lamAxisE, lamAxisI = ts_diag(ts_params, dummy_batch) diff --git a/tests/test_forward/test_angular_1d.py b/tests/test_forward/test_angular_1d.py index 8b6ba66a..b5bb42b2 100644 --- a/tests/test_forward/test_angular_1d.py +++ b/tests/test_forward/test_angular_1d.py @@ -67,7 +67,7 @@ def test_arts1d_forward_pass(): "i_amps": np.array([1]), } - ts_diag = ThomsonScatteringDiagnostic(config, scattering_angles=None, angular=True) + ts_diag = ThomsonScatteringDiagnostic(config, angular=True, cumulative=False) config = ts_diag.get_cfg() ts_params = ThomsonParams(config["parameters"], num_params=1, batch=False, activate=True) ThryE, ThryI, lamAxisE, lamAxisI = ts_diag(ts_params, dummy_batch) diff --git a/tests/test_forward/test_angular_2d.py b/tests/test_forward/test_angular_2d.py index 47bda087..6fa65072 100644 --- a/tests/test_forward/test_angular_2d.py +++ b/tests/test_forward/test_angular_2d.py @@ -77,7 +77,7 @@ def test_arts2d_forward_pass(): "i_amps": np.array([1]), } - ts_diag = ThomsonScatteringDiagnostic(config, scattering_angles=sas) + ts_diag = ThomsonScatteringDiagnostic(config, angular=True, cumulative=False) ts_params = ThomsonParams(config["parameters"], num_params=1, batch=False) ThryE, ThryI, lamAxisE, lamAxisI = ts_diag(ts_params, dummy_batch) # np.save("tests/test_forward/ThryE-arts2d.npy", ThryE) diff --git a/tests/test_inverse/test_1d_random.py b/tests/test_inverse/test_1d_random.py index ea15b8a7..2b1aebeb 100644 --- a/tests/test_inverse/test_1d_random.py +++ b/tests/test_inverse/test_1d_random.py @@ -100,7 +100,7 @@ def test_1d_inverse(): "i_amps": np.array([1]), } rng = np.random.default_rng() - ts_diag = ThomsonScatteringDiagnostic(config, scattering_angles=None, angular=False) + ts_diag = ThomsonScatteringDiagnostic(config, angular=False, cumulative=False) config = ts_diag.get_cfg() config["parameters"] = _perturb_params_(rng, config["parameters"]) misc.log_mlflow(config) @@ -110,7 +110,7 @@ def test_1d_inverse(): loss = 1 while np.nan_to_num(loss, nan=1) > 1e-3: - ts_diag = ThomsonScatteringDiagnostic(config, scattering_angles=None, angular=False) + ts_diag = ThomsonScatteringDiagnostic(config, angular=False, cumulative=False) config["parameters"] = _perturb_params_(rng, config["parameters"]) ts_params_fit = ThomsonParams(config["parameters"], num_params=1, batch=True, activate=True) diff_params, static_params = eqx.partition( diff --git a/tsadar/core/thomson_diagnostic.py b/tsadar/core/thomson_diagnostic.py index 9a3ed247..7f32aa29 100644 --- a/tsadar/core/thomson_diagnostic.py +++ b/tsadar/core/thomson_diagnostic.py @@ -1,6 +1,9 @@ +import os + from jax import numpy as jnp, vmap from tsadar.utils.data_handling.calibration import get_scattering_angles, get_calibrations +from tsadar.utils.data_handling.load_ts_data import loadData from .modules import ThomsonParams from .physics import irf @@ -8,6 +11,8 @@ + + class ThomsonScatteringDiagnostic: """ The SpectrumCalculator class wraps the FitModel class adding instrumental effects to the calculated spectrum so it @@ -22,11 +27,11 @@ class ThomsonScatteringDiagnostic: weights of each of the scattering angles in the final spectrum """ - def __init__(self, cfg, angular=False): + def __init__(self, cfg, angular=False, cumulative=True): super().__init__() - self.cfg, self.scattering_angles = self.initialize_scattering_angles(cfg, angular) + self.cfg, self.scattering_angles = self.initialize_scattering_angles(cfg, angular, cumulative) self.model = FitModel(cfg, self.scattering_angles) @@ -130,7 +135,7 @@ def __call__(self, ts_params: ThomsonParams, batch): return ThryE, ThryI, lamAxisE, lamAxisI - def initialize_scattering_angles(self, config, angular): + def initialize_scattering_angles(self, config, angular, cumulative): """ Initializes scattering angles and weights. @@ -140,24 +145,39 @@ def initialize_scattering_angles(self, config, angular): Returns: Updated configuration dictionary with scattering angles and weights """ - config["other"]["lamrangE"] = [ - config["data"]["fit_rng"]["forward_epw_start"], - config["data"]["fit_rng"]["forward_epw_end"], - ] - config["other"]["lamrangI"] = [ - config["data"]["fit_rng"]["forward_iaw_start"], - config["data"]["fit_rng"]["forward_iaw_end"], - ] - config["other"]["npts"] = int(config["other"]["CCDsize"][1] * config["other"]["points_per_pixel"]) - scattering_angles = get_scattering_angles(config) - - if angular: - [axisxE, _, _, _, _, _] = get_calibrations( - 104000, config["other"]["extraoptions"]["spectype"], 0.0, config["other"]["CCDsize"] - ) # shot number hardcoded to get calibration - config["other"]["extraoptions"]["spectype"] = "angular_full" - - scattering_angles["angAxis"] = axisxE + if cumulative: + custom_path = None + if "filenames" in config["data"].keys(): + if config["data"]["filenames"]["epw"] is not None: + custom_path = os.path.dirname(config["data"]["filenames"]["epw-local"]) + + if config["data"]["filenames"]["iaw"] is not None: + custom_path = os.path.dirname(config["data"]["filenames"]["iaw-local"]) + + [elecData, ionData, xlab, t0, config["other"]["extraoptions"]["spectype"]] = loadData( + config["data"]["shotnum"], config["data"]["shotDay"], config["other"]["extraoptions"], custom_path=custom_path + ) + scattering_angles = get_scattering_angles(config) + + else: + config["other"]["lamrangE"] = [ + config["data"]["fit_rng"]["forward_epw_start"], + config["data"]["fit_rng"]["forward_epw_end"], + ] + config["other"]["lamrangI"] = [ + config["data"]["fit_rng"]["forward_iaw_start"], + config["data"]["fit_rng"]["forward_iaw_end"], + ] + config["other"]["npts"] = int(config["other"]["CCDsize"][1] * config["other"]["points_per_pixel"]) + scattering_angles = get_scattering_angles(config) + + if angular: + [axisxE, _, _, _, _, _] = get_calibrations( + 104000, config["other"]["extraoptions"]["spectype"], 0.0, config["other"]["CCDsize"] + ) # shot number hardcoded to get calibration + config["other"]["extraoptions"]["spectype"] = "angular_full" + + scattering_angles["angAxis"] = axisxE return config, scattering_angles From 60cb2994c9d236bad7b489c2754fda1c5c23165d Mon Sep 17 00:00:00 2001 From: rohchav Date: Mon, 24 Feb 2025 13:49:57 -0600 Subject: [PATCH 3/4] Remove deprecated scattering angle calculations and related test code --- test_1d.py | 61 --------------------------- tests/test_forward/test_1d.py | 11 ----- tests/test_forward/test_angular_1d.py | 19 --------- tests/test_inverse/test_1d_random.py | 12 ------ 4 files changed, 103 deletions(-) delete mode 100644 test_1d.py diff --git a/test_1d.py b/test_1d.py deleted file mode 100644 index e982df64..00000000 --- a/test_1d.py +++ /dev/null @@ -1,61 +0,0 @@ -import time, pytest -import multiprocessing as mp -import yaml -import mlflow -from flatten_dict import flatten, unflatten -from numpy.testing import assert_allclose -from jax import config - -config.update("jax_enable_x64", True) -# config.update("jax_disable_jit", True) -# config.update("jax_check_tracer_leaks", True) - -from tsadar.inverse import fitter -from tsadar.utils import misc - - -@pytest.mark.parametrize("nn", [False]) -def test_data(nn): - # Test #3: Data test, compare fit to a preknown fit result - # currently just runs one line of shot 101675 for the electron, should be expanded in the future - - with open("tests/configs/time_test_defaults.yaml", "r") as fi: - defaults = yaml.safe_load(fi) - - with open("tests/configs/time_test_inputs.yaml", "r") as fi: - inputs = yaml.safe_load(fi) - - defaults = flatten(defaults) - defaults.update(flatten(inputs)) - config = unflatten(defaults) - - config["nn"]["use"] = nn - # config["parameters"]["Te"]["val"] = 0.5 - # config["parameters"]["ne"]["val"] = 0.2 # 0.25 - # config["parameters"]["m"]["val"] = 3.0 # 2.2 - - mlflow.set_experiment(config["mlflow"]["experiment"]) - - with mlflow.start_run() as run: - misc.log_mlflow(config) - config["num_cores"] = int(mp.cpu_count()) - - t0 = time.time() - fit_results, loss = fitter.fit(config=config) - metrics_dict = {"total_time": time.time() - t0, "num_cores": int(mp.cpu_count())} - mlflow.log_metrics(metrics=metrics_dict) - mlflow.set_tag("status", "completed") - print(fit_results) - - # These were changed 5/6/24 to reflect new good fit values, unclear why changes were required probably a change - # to the calibration - assert_allclose(fit_results["amp1_general"][0], 0.734, rtol=1e-1) # 0.9257 - assert_allclose(fit_results["amp2_general"][0], 0.519, rtol=1e-1) # 0.6727 - assert_allclose(fit_results["lam_general"][0], 524.016, rtol=5e-3) # 524.2455 - assert_allclose(fit_results["Te_electron"][0], 0.5994, rtol=1e-1) # 0.67585 - assert_allclose(fit_results["ne_electron"][0], 0.2256, rtol=5e-2) # 0.21792 - assert_allclose(fit_results["m_electron"][0], 2.987, rtol=15e-2) # 3.3673 - - -if __name__ == "__main__": - test_data(False) diff --git a/tests/test_forward/test_1d.py b/tests/test_forward/test_1d.py index c4d2198d..42ce3edb 100644 --- a/tests/test_forward/test_1d.py +++ b/tests/test_forward/test_1d.py @@ -39,17 +39,6 @@ def test_1d_forward_pass(): defaults.update(flatten(inputs)) config = unflatten(defaults) - # get scattering angles and weights - # config["other"]["lamrangE"] = [ - # config["data"]["fit_rng"]["forward_epw_start"], - # config["data"]["fit_rng"]["forward_epw_end"], - # ] - # config["other"]["lamrangI"] = [ - # config["data"]["fit_rng"]["forward_iaw_start"], - # config["data"]["fit_rng"]["forward_iaw_end"], - # ] - # config["other"]["npts"] = int(config["other"]["CCDsize"][1] * config["other"]["points_per_pixel"]) - # sas = get_scattering_angles(config) dummy_batch = { "i_data": np.array([1]), diff --git a/tests/test_forward/test_angular_1d.py b/tests/test_forward/test_angular_1d.py index b5bb42b2..f3c0e7ec 100644 --- a/tests/test_forward/test_angular_1d.py +++ b/tests/test_forward/test_angular_1d.py @@ -39,25 +39,6 @@ def test_arts1d_forward_pass(): defaults.update(flatten(inputs)) config = unflatten(defaults) - # get scattering angles and weights - # config["other"]["lamrangE"] = [ - # config["data"]["fit_rng"]["forward_epw_start"], - # config["data"]["fit_rng"]["forward_epw_end"], - # ] - # config["other"]["lamrangI"] = [ - # config["data"]["fit_rng"]["forward_iaw_start"], - # config["data"]["fit_rng"]["forward_iaw_end"], - # ] - # config["other"]["npts"] = int(config["other"]["CCDsize"][1] * config["other"]["points_per_pixel"]) - # sas = get_scattering_angles(config) - - # [axisxE, _, _, _, _, _] = get_calibrations( - # 104000, config["other"]["extraoptions"]["spectype"], 0.0, config["other"]["CCDsize"] - # ) # shot number hardcoded to get calibration - # config["other"]["extraoptions"]["spectype"] = "angular_full" - - # sas["angAxis"] = axisxE - dummy_batch = { "i_data": np.ones((config["other"]["CCDsize"][0], config["other"]["CCDsize"][1])), "e_data": np.ones((config["other"]["CCDsize"][0], config["other"]["CCDsize"][1])), diff --git a/tests/test_inverse/test_1d_random.py b/tests/test_inverse/test_1d_random.py index 2b1aebeb..43924527 100644 --- a/tests/test_inverse/test_1d_random.py +++ b/tests/test_inverse/test_1d_random.py @@ -79,18 +79,6 @@ def test_1d_inverse(): defaults.update(flatten(inputs)) config = unflatten(defaults) - # get scattering angles and weights - # config["other"]["lamrangE"] = [ - # config["data"]["fit_rng"]["forward_epw_start"], - # config["data"]["fit_rng"]["forward_epw_end"], - # ] - # config["other"]["lamrangI"] = [ - # config["data"]["fit_rng"]["forward_iaw_start"], - # config["data"]["fit_rng"]["forward_iaw_end"], - # ] - # config["other"]["npts"] = int(config["other"]["CCDsize"][1] * config["other"]["points_per_pixel"]) - # sas = get_scattering_angles(config) - dummy_batch = { "i_data": np.array([1]), "e_data": np.array([1]), From 7d33264b1e4a60714a30df747ae76762ebb10e92 Mon Sep 17 00:00:00 2001 From: rohchav Date: Mon, 24 Feb 2025 13:58:47 -0600 Subject: [PATCH 4/4] Remove commented-out scattering angle imports from test files --- tests/test_forward/test_1d.py | 1 - tests/test_forward/test_angular_1d.py | 2 -- tests/test_inverse/test_1d_random.py | 1 - 3 files changed, 4 deletions(-) diff --git a/tests/test_forward/test_1d.py b/tests/test_forward/test_1d.py index 42ce3edb..6de3c85a 100644 --- a/tests/test_forward/test_1d.py +++ b/tests/test_forward/test_1d.py @@ -11,7 +11,6 @@ from tsadar.utils import misc from tsadar.core.thomson_diagnostic import ThomsonScatteringDiagnostic from tsadar.core.modules import ThomsonParams -# from tsadar.utils.data_handling.calibration import get_scattering_angles def test_1d_forward_pass(): diff --git a/tests/test_forward/test_angular_1d.py b/tests/test_forward/test_angular_1d.py index f3c0e7ec..3696ee07 100644 --- a/tests/test_forward/test_angular_1d.py +++ b/tests/test_forward/test_angular_1d.py @@ -11,8 +11,6 @@ from tsadar.core.thomson_diagnostic import ThomsonScatteringDiagnostic from tsadar.core.modules import ThomsonParams -# from tsadar.utils.data_handling.calibration import get_scattering_angles, get_calibrations - def test_arts1d_forward_pass(): """ diff --git a/tests/test_inverse/test_1d_random.py b/tests/test_inverse/test_1d_random.py index 43924527..fe2d647d 100644 --- a/tests/test_inverse/test_1d_random.py +++ b/tests/test_inverse/test_1d_random.py @@ -14,7 +14,6 @@ from tsadar.utils import misc from tsadar.core.thomson_diagnostic import ThomsonScatteringDiagnostic from tsadar.core.modules import ThomsonParams, get_filter_spec -# from tsadar.utils.data_handling.calibration import get_scattering_angles def _perturb_params_(rng, params):