From b7425f1a7738b413f9869f54e35b1537cb583b3d Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 28 Mar 2023 14:55:23 +0100 Subject: [PATCH 001/160] initial commit --- indica/models/magnetic_recon.py | 87 +++++++++++++++++++++++++++++++++ 1 file changed, 87 insertions(+) create mode 100644 indica/models/magnetic_recon.py diff --git a/indica/models/magnetic_recon.py b/indica/models/magnetic_recon.py new file mode 100644 index 00000000..1cc3667a --- /dev/null +++ b/indica/models/magnetic_recon.py @@ -0,0 +1,87 @@ +import xarray as xr +import matplotlib.pyplot as plt + +from indica.models.abstractdiagnostic import DiagnosticModel +from indica.models.plasma import example_run as example_plasma +from indica.readers.available_quantities import AVAILABLE_QUANTITIES + + +class MagneticRecon(DiagnosticModel): + """ + Object representing observations from a magnetic reconstruction + """ + + def __init__( + self, + name: str, + instrument_method="get_equilibrium", + ): + self.name = name + self.instrument_method = instrument_method + self.quantities = AVAILABLE_QUANTITIES[self.instrument_method] + + def _build_bckc_dictionary(self): + self.bckc = {} + + for quant in self.quantities: + datatype = self.quantities[quant] + if quant == "wp": + self.bckc[quant] = self.wp + error = xr.full_like(self.bckc[quant], 0.0) + stdev = xr.full_like(self.bckc[quant], 0.0) + self.bckc[quant].attrs = { + "datatype": datatype, + "error": error, + "stdev": stdev, + "provenance": str(self), + } + else: + # print(f"{quant} not available in model for {self.instrument_method}") + continue + + def __call__( + self, + t = None, + **kwargs, + ): + """ + + Returns + ------- + bckc values + """ + if self.plasma is None: + raise ValueError("plasma object is needed") + + if t is None: + t = self.plasma.time_to_calculate + + self.wp = self.plasma.wp.sel(t=t) + self._build_bckc_dictionary() + return self.bckc + + +def example_run( + diagnostic_name: str = "efit", + plasma=None, + plot=False, + t=None, +): + if plasma is None: + plasma = example_plasma() + + model = MagneticRecon( + diagnostic_name, + ) + model.set_plasma(plasma) + bckc = model() + + if plot: + bckc["wp"].plot() + + return plasma, model, bckc + + +if __name__ == "__main__": + example_run(plot=True) + plt.show(block=True) From aae4c4726ecd1c0d9cd28a8393a7e78079f98f9e Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 28 Mar 2023 14:55:46 +0100 Subject: [PATCH 002/160] magnetic recon tests --- tests/unit/models/test_diagnostic_models.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/tests/unit/models/test_diagnostic_models.py b/tests/unit/models/test_diagnostic_models.py index 5055a93a..00323ab8 100644 --- a/tests/unit/models/test_diagnostic_models.py +++ b/tests/unit/models/test_diagnostic_models.py @@ -11,6 +11,8 @@ from indica.models.interferometry import example_run as interf from indica.models.plasma import example_run as example_plasma from indica.models.thomson_scattering import example_run as ts +from indica.models.magnetic_recon import example_run as mag_recon + EXAMPLES: Dict[str, Callable] = { "bolometer_camera": bolo, @@ -19,6 +21,7 @@ "helike_spectroscopy": helike, "thomson_scattering": ts, "charge_exchange": cxrs, + "magnetic_reconstruction":mag_recon } PLASMA = example_plasma(pulse=None, tstart=0, tend=0.1, dt=0.02) NT = np.size(PLASMA.t) @@ -124,3 +127,15 @@ def test_helike_timepoint_pass(): def test_helike_interpolation(): _test_time_interpolation("helike_spectroscopy") + + +def test_magrecon_timepoint_fail(): + _test_timepoint_fail("magnetic_reconstruction") + + +def test_magrecon_timepoint_pass(): + _test_timepoint_pass("magnetic_reconstruction") + + +def test_magrecon_interpolation(): + _test_time_interpolation("magnetic_reconstruction") \ No newline at end of file From 3100262d7702c7917923a510d5f13c90c7b25535 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 30 Mar 2023 10:06:18 +0100 Subject: [PATCH 003/160] use error attribute if it is given --- indica/bayesmodels.py | 34 ++++++++++++++++++++-------------- 1 file changed, 20 insertions(+), 14 deletions(-) diff --git a/indica/bayesmodels.py b/indica/bayesmodels.py index bba9d892..6a4c2ef4 100644 --- a/indica/bayesmodels.py +++ b/indica/bayesmodels.py @@ -59,7 +59,17 @@ def __init__( raise ValueError(f"{missing_data} not found in data given") def _build_bckc(self, params: dict, **kwargs): - # TODO: consider how to handle if models have overlapping kwargs + """ + TODO: consider how to handle if models have overlapping kwargs + Parameters + ---------- + params - dictionary which is updated by optimiser + kwargs - passed to model i.e. settings + + Returns + ------- + bckc of results + """ # Params is a dictionary which is updated by optimiser, # kwargs is constant i.e. settings for models self.bckc: dict = {} @@ -73,22 +83,18 @@ def _build_bckc(self, params: dict, **kwargs): def _ln_likelihood(self): ln_likelihood = 0 for key in self.quant_to_optimise: - # TODO: What to use as error? Assume percentage error if none given... # Float128 is used since rounding of small numbers causes # problems when initial results are bad fits model_data = self.bckc[key].values.astype("float128") - exp_data = ( - self.data[key] - .sel(t=self.plasma.time_to_calculate) - .values.astype("float128") - ) - _ln_likelihood = np.log( - gaussian( - model_data, - exp_data, - exp_data * 0.10 - ) - ) + exp_data = self.data[key].sel(t=self.plasma.time_to_calculate).values.astype("float128") + if hasattr(self.data[key], "error"): # Assume percentage error if none given. + if (self.data[key].error != 0).any(): # TODO: Some models have an error of 0 given + exp_error = self.data[key].error.values + else: + exp_error = exp_data * 0.10 + else: + exp_error = exp_data * 0.10 + _ln_likelihood = np.log(gaussian(model_data, exp_data, exp_error)) ln_likelihood += np.nanmean(_ln_likelihood) return ln_likelihood From ba574ecc804e2d2a1518fbf1b8c07a58965e5ae4 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 30 Mar 2023 10:06:56 +0100 Subject: [PATCH 004/160] Added poisson noise and adjusted calibration factor --- indica/models/helike_spectroscopy.py | 132 ++++++++++++++------------- 1 file changed, 68 insertions(+), 64 deletions(-) diff --git a/indica/models/helike_spectroscopy.py b/indica/models/helike_spectroscopy.py index e5a1cd4f..311aeea7 100644 --- a/indica/models/helike_spectroscopy.py +++ b/indica/models/helike_spectroscopy.py @@ -36,17 +36,17 @@ class Helike_spectroscopy(DiagnosticModel): """ def __init__( - self, - name: str, - instrument_method="get_helike_spectroscopy", - etendue: float = 1.0, - calibration: float = 1.0e-27, - marchuk: bool = True, - full_run: bool = False, - element: str = "ar", - window_len: int = 1030, - window_lim: list = [0.394, 0.401], - window_masks: list = [], + self, + name: str, + instrument_method="get_helike_spectroscopy", + etendue: float = 1.0, + calibration: float = 5.0e-18, + marchuk: bool = True, + full_run: bool = False, + element: str = "ar", + window_len: int = 1030, + window_lim: list = [0.394, 0.401], + window_masks: list = [], ): """ Read all atomic data and initialise objects @@ -159,8 +159,8 @@ def _set_adas_pecs(self): for line in pec: pec[line]["emiss_coeff"] = ( pec[line]["emiss_coeff"] - .sel(electron_density=4.0e19, method="nearest") - .drop_vars("electron_density") + .sel(electron_density=4.0e19, method="nearest") + .drop_vars("electron_density") ) self.pec = pec @@ -231,27 +231,27 @@ def _calculate_line_emission(self, line_labels=["w", "k"]): ) for line_name in line_labels: _line_emission = ( - _pecs.isel( - line_name=(_pecs.wavelength < self.line_ranges[line_name].stop) - & (_pecs.wavelength > self.line_ranges[line_name].start) - ) - * mult + _pecs.isel( + line_name=(_pecs.wavelength < self.line_ranges[line_name].stop) + & (_pecs.wavelength > self.line_ranges[line_name].start) + ) + * mult ) _line_emission = _line_emission.sum(["type", "line_name"]) # TODO: convert PEC wavelengths to nm as per convention at TE! ev_wavelength = ph.nm_eV_conversion(nm=self.line_ranges[line_name].start) line_emission[line_name] = ( - xr.where(_line_emission >= 0, _line_emission, 0) * ev_wavelength - ) * self.calibration + xr.where(_line_emission >= 0, _line_emission, 0) * ev_wavelength + ) * self.calibration if "k" in line_emission.keys() and "w" in line_emission.keys(): line_emission["kw"] = line_emission["k"] * line_emission["w"] if "n3" in line_emission.keys() and "w" in line_emission.keys(): line_emission["n3w"] = line_emission["n345"] * line_emission["w"] if ( - "n3" in line_emission.keys() - and "n345" in line_emission.keys() - and "w" in line_emission.keys() + "n3" in line_emission.keys() + and "n345" in line_emission.keys() + and "w" in line_emission.keys() ): line_emission["tot"] = line_emission["n345"] + line_emission["w"] line_emission["n3tot"] = line_emission["n345"] * line_emission["w"] @@ -280,6 +280,8 @@ def _calculate_los_integral(self, calc_rho=False): t=self.spectra.t, calc_rho=calc_rho, ) + # TODO: LOS integral removes NaNs so manually add them back (find better solution) + self.measured_spectra[self.measured_spectra == 0] = np.nan def _calculate_temperatures(self): x1 = self.los_transform.x1 @@ -317,10 +319,10 @@ def _calculate_temperatures(self): ) def _moment_analysis( - self, - line: str, - profile_1d: DataArray = None, - half_los: bool = True, + self, + line: str, + profile_1d: DataArray = None, + half_los: bool = True, ): """ Perform moment analysis using a specific line emission as distribution function @@ -453,7 +455,7 @@ def _moment_analysis( return result, pos, err_in, err_out def _calculate_intensity( - self, + self, ): """ Returns DataArrays of emission type with co-ordinates of line label and @@ -467,8 +469,8 @@ def _calculate_intensity( _pecs_ds = _pecs.to_dataset("type") temp = [ _pecs_ds[type] - .dropna("line_name", how="all") - .interp(electron_temperature=self.Te, assume_sorted=True) + .dropna("line_name", how="all") + .interp(electron_temperature=self.Te, assume_sorted=True) for type in _pecs_ds.data_vars.keys() ] _intensity = xr.merge(temp).to_array("type") @@ -476,7 +478,7 @@ def _calculate_intensity( return intensity def _make_spectra( - self, + self, ): """ TODO: Doppler Shift / Add convolution of broadening @@ -525,29 +527,29 @@ def _build_bckc_dictionary(self): datatype = self.quantities[quant] if datatype[0] == ("intensity"): line = str(quant.split("_")[1]) - quantity = f"int_{line}" - self.bckc[quantity] = self.measured_intensity[line] - self.bckc[quantity].attrs["emiss"] = self.line_emission[line] + quant = f"int_{line}" + self.bckc[quant] = self.measured_intensity[line] + self.bckc[quant].attrs["emiss"] = self.line_emission[line] elif datatype == ("temperature", "electrons"): line = str(quant.split("_")[1]) - quantity = f"te_{line}" - self.bckc[quantity] = self.measured_Te[line] - self.bckc[quantity].attrs["emiss"] = self.line_emission[line] + quant = f"te_{line}" + self.bckc[quant] = self.measured_Te[line] + self.bckc[quant].attrs["emiss"] = self.line_emission[line] elif datatype == ("temperature", "ions"): line = str(quant.split("_")[1]) - quantity = f"ti_{line}" - self.bckc[quantity] = self.measured_Ti[line] - self.bckc[quantity].attrs["emiss"] = self.line_emission[line] + quant = f"ti_{line}" + self.bckc[quant] = self.measured_Ti[line] + self.bckc[quant].attrs["emiss"] = self.line_emission[line] elif datatype == ("spectra", "passive"): - quantity = quant self.bckc["spectra"] = self.measured_spectra + self.bckc["spectra"].attrs["error"] = np.sqrt(self.measured_spectra) # poisson noise else: print(f"{quant} not available in model for {self.instrument_method}") continue - self.bckc[quantity].attrs["datatype"] = datatype + self.bckc[quant].attrs["datatype"] = datatype if quant != "spectra" and hasattr(self, "spectra"): - self.bckc[quantity].attrs["pos"] = { + self.bckc[quant].attrs["pos"] = { "value": self.pos[line], "err_in": self.err_in[line], "err_out": self.err_out[line], @@ -560,19 +562,19 @@ def _build_bckc_dictionary(self): self.bckc["int_n3/int_tot"] = self.bckc["int_n3"] / self.bckc["int_tot"] def __call__( - self, - Te: DataArray = None, - Ti: DataArray = None, - Ne: DataArray = None, - Nimp: DataArray = None, - Fz: dict = None, - Nh: DataArray = None, - t: LabeledArray = None, - calc_spectra=True, - calc_rho: bool = False, - minimum_lines: bool = False, - moment_analysis: bool = False, - **kwargs, + self, + Te: DataArray = None, + Ti: DataArray = None, + Ne: DataArray = None, + Nimp: DataArray = None, + Fz: dict = None, + Nh: DataArray = None, + t: LabeledArray = None, + calc_spectra=True, + calc_rho: bool = False, + minimum_lines: bool = False, + moment_analysis: bool = False, + **kwargs, ): """ Calculate diagnostic measured values @@ -622,12 +624,12 @@ def __call__( Nimp = self.plasma.impurity_density.sel(t=t) else: if ( - Ne is None - or Te is None - or Nh is None - or Fz is None - or Ti is None - or Nimp is None + Ne is None + or Te is None + or Nh is None + or Fz is None + or Ti is None + or Nimp is None ): raise ValueError("Give inputs or assign plasma class!") @@ -649,7 +651,7 @@ def __call__( if minimum_lines: line_labels = [ "w", - "k", + "k" ] names = ["int_w", "int_k", "te_kw", "ti_w"] else: @@ -662,7 +664,7 @@ def __call__( "te_n3w", "te_kw", "ti_z", - "ti_w", + "ti_w" ] quant = dict(quant, **{x: self.quantities[x] for x in names}) @@ -700,6 +702,8 @@ def doppler_broaden(x, integral, center, ion_mass, ion_temp): def gaussian(x, integral, center, sigma): return (integral / (sigma * np.sqrt(2 * np.pi)) * np.exp(-((x - center) ** 2) / (2 * sigma ** 2))) + + # fmt: on From d5fec79bbbd25a34805422dcd82c2f78e038596f Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 30 Mar 2023 10:41:05 +0100 Subject: [PATCH 005/160] changed name magnetic_recon -> equilibrium_reconstruction --- ...ic_recon.py => equilibrium_reconstruction.py} | 6 +++--- tests/unit/models/test_diagnostic_models.py | 16 ++++++++-------- 2 files changed, 11 insertions(+), 11 deletions(-) rename indica/models/{magnetic_recon.py => equilibrium_reconstruction.py} (94%) diff --git a/indica/models/magnetic_recon.py b/indica/models/equilibrium_reconstruction.py similarity index 94% rename from indica/models/magnetic_recon.py rename to indica/models/equilibrium_reconstruction.py index 1cc3667a..95dff261 100644 --- a/indica/models/magnetic_recon.py +++ b/indica/models/equilibrium_reconstruction.py @@ -6,7 +6,7 @@ from indica.readers.available_quantities import AVAILABLE_QUANTITIES -class MagneticRecon(DiagnosticModel): +class EquilibriumReconstruction(DiagnosticModel): """ Object representing observations from a magnetic reconstruction """ @@ -70,8 +70,8 @@ def example_run( if plasma is None: plasma = example_plasma() - model = MagneticRecon( - diagnostic_name, + model = EquilibriumReconstruction( + diagnostic_name ) model.set_plasma(plasma) bckc = model() diff --git a/tests/unit/models/test_diagnostic_models.py b/tests/unit/models/test_diagnostic_models.py index 00323ab8..2e01a538 100644 --- a/tests/unit/models/test_diagnostic_models.py +++ b/tests/unit/models/test_diagnostic_models.py @@ -11,7 +11,7 @@ from indica.models.interferometry import example_run as interf from indica.models.plasma import example_run as example_plasma from indica.models.thomson_scattering import example_run as ts -from indica.models.magnetic_recon import example_run as mag_recon +from indica.models.equilibrium_reconstruction import example_run as equil_recon EXAMPLES: Dict[str, Callable] = { @@ -21,7 +21,7 @@ "helike_spectroscopy": helike, "thomson_scattering": ts, "charge_exchange": cxrs, - "magnetic_reconstruction":mag_recon + "equilibrium_reconstruction": equil_recon, } PLASMA = example_plasma(pulse=None, tstart=0, tend=0.1, dt=0.02) NT = np.size(PLASMA.t) @@ -129,13 +129,13 @@ def test_helike_interpolation(): _test_time_interpolation("helike_spectroscopy") -def test_magrecon_timepoint_fail(): - _test_timepoint_fail("magnetic_reconstruction") +def test_equil_recon_timepoint_fail(): + _test_timepoint_fail("equilibrium_reconstruction") -def test_magrecon_timepoint_pass(): - _test_timepoint_pass("magnetic_reconstruction") +def test_equil_recon_timepoint_pass(): + _test_timepoint_pass("equilibrium_reconstruction") -def test_magrecon_interpolation(): - _test_time_interpolation("magnetic_reconstruction") \ No newline at end of file +def test_equil_recon_interpolation(): + _test_time_interpolation("equilibrium_reconstruction") \ No newline at end of file From 56c769a83929eb76ff0f92e0a7b6340e276685fa Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 30 Mar 2023 15:31:35 +0100 Subject: [PATCH 006/160] diagnostic data with channel dim are treated as individual measurements --- indica/bayesmodels.py | 27 ++++++++++++++------------- 1 file changed, 14 insertions(+), 13 deletions(-) diff --git a/indica/bayesmodels.py b/indica/bayesmodels.py index 6a4c2ef4..c99263c6 100644 --- a/indica/bayesmodels.py +++ b/indica/bayesmodels.py @@ -55,12 +55,12 @@ def __init__( diag_model.plasma = self.plasma missing_data = list(set(quant_to_optimise).difference(data.keys())) - if missing_data: # list of keys in quant_to_optimise but not data + if missing_data: # gives list of keys in quant_to_optimise but not data raise ValueError(f"{missing_data} not found in data given") def _build_bckc(self, params: dict, **kwargs): """ - TODO: consider how to handle if models have overlapping kwargs + TODO: consider how to handle when models have overlapping kwargs Parameters ---------- params - dictionary which is updated by optimiser @@ -83,19 +83,20 @@ def _build_bckc(self, params: dict, **kwargs): def _ln_likelihood(self): ln_likelihood = 0 for key in self.quant_to_optimise: - # Float128 is used since rounding of small numbers causes - # problems when initial results are bad fits - model_data = self.bckc[key].values.astype("float128") - exp_data = self.data[key].sel(t=self.plasma.time_to_calculate).values.astype("float128") - if hasattr(self.data[key], "error"): # Assume percentage error if none given. + # Float128 since rounding of small numbers causes problems when initial results are bad fits + model_data = self.bckc[key].astype("float128") + exp_data = self.data[key].sel(t=self.plasma.time_to_calculate).astype("float128") + + exp_error = exp_data * 0.10 # Assume percentage error if none given. + if hasattr(self.data[key], "error"): if (self.data[key].error != 0).any(): # TODO: Some models have an error of 0 given - exp_error = self.data[key].error.values - else: - exp_error = exp_data * 0.10 - else: - exp_error = exp_data * 0.10 + exp_error = self.data[key].error + _ln_likelihood = np.log(gaussian(model_data, exp_data, exp_error)) - ln_likelihood += np.nanmean(_ln_likelihood) + # Treat channel as key dim which isn't averaged like other dims + if "channel" in _ln_likelihood.dims: + _ln_likelihood = _ln_likelihood.sum(dim="channel", skipna=True) + ln_likelihood += _ln_likelihood.mean(skipna=True).values return ln_likelihood def _ln_prior(self, parameters: dict): From 91d7e66d08565b2f1709090f62973b6e95ac6b44 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Fri, 31 Mar 2023 13:22:59 +0100 Subject: [PATCH 007/160] added function of example LOS --- indica/models/helike_spectroscopy.py | 39 ++++++++++--------- .../models/test_diagnostic_models.py | 0 .../models/test_plasma.py | 0 3 files changed, 21 insertions(+), 18 deletions(-) rename tests/{unit => integration}/models/test_diagnostic_models.py (100%) rename tests/{unit => integration}/models/test_plasma.py (100%) diff --git a/indica/models/helike_spectroscopy.py b/indica/models/helike_spectroscopy.py index e5a1cd4f..175ee401 100644 --- a/indica/models/helike_spectroscopy.py +++ b/indica/models/helike_spectroscopy.py @@ -742,16 +742,7 @@ def select_transition(adf15_data, transition: str, wavelength: float): return pec -def example_run(pulse: int = 9229, plasma=None, plot=False, **kwargs): - # TODO: LOS sometimes crossing bad EFIT reconstruction - if plasma is None: - plasma = example_plasma( - pulse=pulse, impurities=("ar",), impurity_concentration=(0.001,), n_rad=10 - ) - plasma.time_to_calculate = plasma.t[5] - # Create new diagnostic - diagnostic_name = "xrcs" - nchannels = 3 +def helike_LOS_example(nchannels=3): los_end = np.full((nchannels, 3), 0.0) los_end[:, 0] = 0.17 los_end[:, 1] = 0.0 @@ -761,16 +752,28 @@ def example_run(pulse: int = 9229, plasma=None, plot=False, **kwargs): direction = los_end - los_start los_transform = LineOfSightTransform( - origin[:, 0], - origin[:, 1], - origin[:, 2], - direction[:, 0], - direction[:, 1], - direction[:, 2], - name=diagnostic_name, - machine_dimensions=plasma.machine_dimensions, + origin[0:nchannels, 0], + origin[0:nchannels, 1], + origin[0:nchannels, 2], + direction[0:nchannels, 0], + direction[0:nchannels, 1], + direction[0:nchannels, 2], + name="diagnostic_name", + machine_dimensions=((0.15, 0.95), (-0.7, 0.7)), passes=1, ) + return los_transform + +def example_run(pulse: int = 9229, plasma=None, plot=False, **kwargs): + # TODO: LOS sometimes crossing bad EFIT reconstruction + if plasma is None: + plasma = example_plasma( + pulse=pulse, impurities=("ar",), impurity_concentration=(0.001,), n_rad=10 + ) + plasma.time_to_calculate = plasma.t[5] + # Create new diagnostic + diagnostic_name = "xrcs" + los_transform = helike_LOS_example(3) los_transform.set_equilibrium(plasma.equilibrium) model = Helike_spectroscopy( diagnostic_name, diff --git a/tests/unit/models/test_diagnostic_models.py b/tests/integration/models/test_diagnostic_models.py similarity index 100% rename from tests/unit/models/test_diagnostic_models.py rename to tests/integration/models/test_diagnostic_models.py diff --git a/tests/unit/models/test_plasma.py b/tests/integration/models/test_plasma.py similarity index 100% rename from tests/unit/models/test_plasma.py rename to tests/integration/models/test_plasma.py From 37cd733a116258f3e7c8115b274e3122ece9ce9c Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Fri, 31 Mar 2023 13:53:56 +0100 Subject: [PATCH 008/160] refactoring --- .../models/test_helike_spectroscopy.py | 83 ++++++------------- 1 file changed, 24 insertions(+), 59 deletions(-) diff --git a/tests/integration/models/test_helike_spectroscopy.py b/tests/integration/models/test_helike_spectroscopy.py index c0ccb061..53ece85f 100644 --- a/tests/integration/models/test_helike_spectroscopy.py +++ b/tests/integration/models/test_helike_spectroscopy.py @@ -1,75 +1,40 @@ -import numpy as np - -from indica.converters.line_of_sight import LineOfSightTransform import indica.models.helike_spectroscopy as helike +from indica.models.helike_spectroscopy import helike_LOS_example from indica.models.plasma import example_run as example_plasma +import pytest -# Setup LOS -def helike_LOS_example(nchannels=3): - los_end = np.full((nchannels, 3), 0.0) - los_end[:, 0] = 0.17 - los_end[:, 1] = 0.0 - los_end[:, 2] = np.linspace(0.43, -0.43, nchannels) - los_start = np.array([[0.8, 0, 0]] * los_end.shape[0]) - origin = los_start - direction = los_end - los_start - - los_transform = LineOfSightTransform( - origin[0:nchannels, 0], - origin[0:nchannels, 1], - origin[0:nchannels, 2], - direction[0:nchannels, 0], - direction[0:nchannels, 1], - direction[0:nchannels, 2], - name="diagnostic_name", - machine_dimensions=((0.15, 0.95), (-0.7, 0.7)), - passes=1, - ) - return los_transform - +plasma = example_plasma(tstart=0, tend=0.1, dt=0.02) class TestHelike: def setup_class(self): - self.plasma = example_plasma(pulse=9229) - self.single_time_point = self.plasma.time_to_calculate[1] - self.multiple_time_point = self.plasma.time_to_calculate + self.single_time_point = plasma.time_to_calculate[1] + self.multiple_time_point = plasma.time_to_calculate self.multiple_channel_los_transform = helike_LOS_example(nchannels=3) self.single_channel_los_transform = helike_LOS_example(nchannels=1) - self.single_channel_los_transform.set_equilibrium(self.plasma.equilibrium) - self.multiple_channel_los_transform.set_equilibrium(self.plasma.equilibrium) + self.single_channel_los_transform.set_equilibrium(plasma.equilibrium) + self.multiple_channel_los_transform.set_equilibrium(plasma.equilibrium) - def test_plasma_time_to_calculate_is_longer_than_one(self): - assert self.plasma.time_to_calculate.__len__() > 1 + def setup_method(self): + self.model = helike.Helike_spectroscopy("diagnostic_name") + self.model.set_plasma(plasma) + + def test_helike_moment_runs_with_multiple_LOS(self, ): + self.model.set_los_transform(self.multiple_channel_los_transform) + bckc = self.model(calc_spectra=False, moment_analysis=True) + assert bckc - def test_helike_runs_with_example_plasma_and_multiple_LOS(self): - model = helike.Helike_spectroscopy( - "diagnostic_name", - ) - self.plasma.time_to_calculate = self.multiple_time_point - model.set_plasma(self.plasma) - model.set_los_transform(self.multiple_channel_los_transform) - bckc = model(calc_spectra=False, moment_analysis=True) + def test_helike_moment_runs_with_single_LOS(self, ): + self.model.set_los_transform(self.single_channel_los_transform) + bckc = self.model(calc_spectra=False, moment_analysis=True) assert bckc - def test_helike_runs_with_example_plasma_and_single_LOS_and_multiple_time_point( - self, - ): - model = helike.Helike_spectroscopy( - "diagnostic_name", - ) - model.set_plasma(self.plasma) - self.plasma.time_to_calculate = self.multiple_time_point - model.set_los_transform(self.single_channel_los_transform) - bckc = model(calc_spectra=False, moment_analysis=True) + def test_helike_spectra_with_multiple_LOS(self, ): + self.model.set_los_transform(self.multiple_channel_los_transform) + bckc = self.model(calc_spectra=True, moment_analysis=False) assert bckc - def test_helike_runs_with_example_plasma_and_single_LOS_and_single_time_point(self): - model = helike.Helike_spectroscopy( - "diagnostic_name", - ) - model.set_plasma(self.plasma) - self.plasma.time_to_calculate = self.single_time_point - model.set_los_transform(self.single_channel_los_transform) - bckc = model(calc_spectra=False, moment_analysis=True) + def test_helike_spectra_with_single_LOS(self, ): + self.model.set_los_transform(self.single_channel_los_transform) + bckc = self.model(calc_spectra=True, moment_analysis=False) assert bckc From 6dbebc64c225071fa644ed4c3807b61967628b15 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Fri, 31 Mar 2023 13:54:09 +0100 Subject: [PATCH 009/160] refactoring --- tests/integration/test_bayesmodels.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/tests/integration/test_bayesmodels.py b/tests/integration/test_bayesmodels.py index 34aaa020..122e43eb 100644 --- a/tests/integration/test_bayesmodels.py +++ b/tests/integration/test_bayesmodels.py @@ -10,7 +10,7 @@ class TestBayesModels: def setup_class(self): - self.plasma = example_run(pulse=9229) + self.plasma = example_run() self.plasma.time_to_calculate = self.plasma.t[1] self.los_transform = helike_LOS_example(nchannels=1) self.los_transform.set_equilibrium(self.plasma.equilibrium) @@ -24,9 +24,7 @@ def test_simple_run_bayesmodels_with_xrcs(self): priors = { "Te_prof.y0": get_uniform(2e3, 5e3), - # "Te_prof_peaking": get_uniform(1, 5), "Ti_prof.y0": get_uniform(2e3, 8e3), - # "Ti_prof_peaking": get_uniform(1, 5), } bckc = {} @@ -47,9 +45,7 @@ def test_simple_run_bayesmodels_with_xrcs(self): # Setup Optimiser param_names = [ "Te_prof.y0", - # "Te_prof.peaking", "Ti_prof.y0", - # "Ti_prof.peaking", ] ndim = param_names.__len__() From 57c6a665e6d9b3994ca1d80841519979a5010afa Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Fri, 31 Mar 2023 13:54:54 +0100 Subject: [PATCH 010/160] refactoring --- tests/integration/test_bayesmodels.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration/test_bayesmodels.py b/tests/integration/test_bayesmodels.py index 122e43eb..72cb7295 100644 --- a/tests/integration/test_bayesmodels.py +++ b/tests/integration/test_bayesmodels.py @@ -1,6 +1,6 @@ import emcee import flatdict -from tests.integration.models.test_helike_spectroscopy import helike_LOS_example +from indica.models.helike_spectroscopy import helike_LOS_example from indica.bayesmodels import BayesModels from indica.bayesmodels import get_uniform From ac904b60b495139ebac03f8ebb11f8b395c70850 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Fri, 31 Mar 2023 13:57:16 +0100 Subject: [PATCH 011/160] initial commit --- tests/unit/models/__init__.py | 0 tests/unit/models/test_plasma.py | 0 2 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/unit/models/__init__.py create mode 100644 tests/unit/models/test_plasma.py diff --git a/tests/unit/models/__init__.py b/tests/unit/models/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/models/test_plasma.py b/tests/unit/models/test_plasma.py new file mode 100644 index 00000000..e69de29b From 4fa0bb1e4948ab3498d8c76fc3196f363aa52495 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Fri, 31 Mar 2023 14:42:33 +0100 Subject: [PATCH 012/160] removed obsolete attributes --- indica/models/plasma.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/indica/models/plasma.py b/indica/models/plasma.py index 5fde7a3b..ec34e0fb 100644 --- a/indica/models/plasma.py +++ b/indica/models/plasma.py @@ -211,10 +211,6 @@ def initialize_variables(self, tstart: float, tend: float, dt: float): self.tend = tend self.dt = dt - # Dictionary keeping track of deta use for optimisations - self.optimisation: dict = {} - self.forward_models: dict = {} - # Assign plasma and machine attributes self.machine_R = np.linspace( self.machine_dimensions[0][0], self.machine_dimensions[0][1], 100 From 4a5c5b3d3f91caae4117bc34c8c54d583763ad77 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Fri, 31 Mar 2023 16:27:11 +0100 Subject: [PATCH 013/160] Setting up testflow --- tests/unit/models/test_plasma.py | 47 ++++++++++++++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/tests/unit/models/test_plasma.py b/tests/unit/models/test_plasma.py index e69de29b..e2a472c0 100644 --- a/tests/unit/models/test_plasma.py +++ b/tests/unit/models/test_plasma.py @@ -0,0 +1,47 @@ +import numpy as np + +from indica.models.plasma import Plasma +from indica.equilibrium import fake_equilibrium_data, Equilibrium + + +class TestPlasmaInit: + + def setup_class(self): + self.tstart = 0. + self.tend = 0.1 + self.dt = 0.02 + self.plasma = Plasma(tstart=self.tstart, tend=self.tend, dt=self.dt, main_ion="h", + impurities=("c", "ar", "he"), impurity_concentration=(0.01, 0.001, 0.01),) + self.equilibrium_data = fake_equilibrium_data( + tstart=self.tstart, tend=self.tend, dt=self.dt, machine_dims=self.plasma.machine_dimensions + ) + self.equilibrium = Equilibrium(self.equilibrium_data) + self.plasma.set_equilibrium(equilibrium=self.equilibrium) + + def teardown_method(self): + self.plasma.initialize_variables(tstart=self.tstart, tend=self.tend, dt=self.dt) + + def test_plasma_initializes(self): + assert hasattr(self, "plasma") + + def test_plasma_volume_non_zero(self): + _volume = self.plasma.volume + assert len(np.where(_volume > 0)[0]) != 0 + + + +# class TestPlasmaProfiles: +# +# def setup_class(self): +# +# def test_update_profiles +# +# def test_assign_profiles +# +# +# +# +# class TestCacheDependency: +# +# def setup_class(self): +# return From 31f306b5fa74026a40126e63d2e010db59cfd0e0 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Fri, 31 Mar 2023 16:48:47 +0100 Subject: [PATCH 014/160] more tests --- tests/unit/models/test_plasma.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/unit/models/test_plasma.py b/tests/unit/models/test_plasma.py index e2a472c0..5786ade1 100644 --- a/tests/unit/models/test_plasma.py +++ b/tests/unit/models/test_plasma.py @@ -18,6 +18,10 @@ def setup_class(self): self.equilibrium = Equilibrium(self.equilibrium_data) self.plasma.set_equilibrium(equilibrium=self.equilibrium) + def setup_method(self): + self.plasma.electron_density = 1 # set profiles + return + def teardown_method(self): self.plasma.initialize_variables(tstart=self.tstart, tend=self.tend, dt=self.dt) @@ -28,7 +32,19 @@ def test_plasma_volume_non_zero(self): _volume = self.plasma.volume assert len(np.where(_volume > 0)[0]) != 0 + def test_fz_is_non_zero(self): + _fz = self.plasma.fz + assert len(np.where(_fz > 0)[0]) != 0 + + def test_lz_is_non_zero(self): + _lz_tot = self.plasma.lz_tot + assert len(np.where(_lz_tot > 0)[0]) != 0 + + def test_fz_one_time_point(self): + return + def test_fz_keys_match_elements(self): + return # class TestPlasmaProfiles: # @@ -45,3 +61,6 @@ def test_plasma_volume_non_zero(self): # # def setup_class(self): # return +if __name__ == "__main__": + test = TestPlasmaInit() + test.setup_class() From 66e6b22f06488c9355b374296c7ff80538e5f785 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Fri, 31 Mar 2023 16:59:05 +0100 Subject: [PATCH 015/160] more test stubs --- tests/unit/models/test_plasma.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/tests/unit/models/test_plasma.py b/tests/unit/models/test_plasma.py index 5786ade1..f1ca96d7 100644 --- a/tests/unit/models/test_plasma.py +++ b/tests/unit/models/test_plasma.py @@ -4,14 +4,16 @@ from indica.equilibrium import fake_equilibrium_data, Equilibrium -class TestPlasmaInit: - +class TestPlasma: def setup_class(self): self.tstart = 0. self.tend = 0.1 self.dt = 0.02 + self.impurities = ("c", "ar", "he") + self.impurity_concentration = (0.01, 0.001, 0.01) + self.plasma = Plasma(tstart=self.tstart, tend=self.tend, dt=self.dt, main_ion="h", - impurities=("c", "ar", "he"), impurity_concentration=(0.01, 0.001, 0.01),) + impurities=self.impurities, impurity_concentration=self.impurity_concentration) self.equilibrium_data = fake_equilibrium_data( tstart=self.tstart, tend=self.tend, dt=self.dt, machine_dims=self.plasma.machine_dimensions ) @@ -28,6 +30,9 @@ def teardown_method(self): def test_plasma_initializes(self): assert hasattr(self, "plasma") + def test_equilibrium_initializes(self): + assert hasattr(self, "equilibrium") + def test_plasma_volume_non_zero(self): _volume = self.plasma.volume assert len(np.where(_volume > 0)[0]) != 0 From b9e61c2d11ee803f9a840dcd7d3da2b90efba7f7 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Wed, 12 Apr 2023 09:20:37 +0100 Subject: [PATCH 016/160] small fix --- tests/unit/models/test_plasma.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/models/test_plasma.py b/tests/unit/models/test_plasma.py index f1ca96d7..6667e367 100644 --- a/tests/unit/models/test_plasma.py +++ b/tests/unit/models/test_plasma.py @@ -67,5 +67,5 @@ def test_fz_keys_match_elements(self): # def setup_class(self): # return if __name__ == "__main__": - test = TestPlasmaInit() + test = TestPlasma() test.setup_class() From ccfe635c1ddd38bfd99aec8fd42613a4f16d73b6 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 25 Apr 2023 14:38:22 +0100 Subject: [PATCH 017/160] cosmetic changes --- indica/models/helike_spectroscopy.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/indica/models/helike_spectroscopy.py b/indica/models/helike_spectroscopy.py index 311aeea7..08ff9fd2 100644 --- a/indica/models/helike_spectroscopy.py +++ b/indica/models/helike_spectroscopy.py @@ -693,8 +693,8 @@ def __call__( # fmt: off def doppler_broaden(x, integral, center, ion_mass, ion_temp): - _mass = (ion_mass * constants.proton_mass * constants.c ** 2) - sigma = (np.sqrt(constants.e / _mass * ion_temp) * center) + mass = (ion_mass * constants.proton_mass * constants.c ** 2) + sigma = np.sqrt(constants.e / mass * ion_temp) * center gaussian_broadened = gaussian(x, integral, center, sigma, ) return gaussian_broadened @@ -702,8 +702,6 @@ def doppler_broaden(x, integral, center, ion_mass, ion_temp): def gaussian(x, integral, center, sigma): return (integral / (sigma * np.sqrt(2 * np.pi)) * np.exp(-((x - center) ** 2) / (2 * sigma ** 2))) - - # fmt: on From c8cd11be44d614905a0ad54cb740c08dace37663 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 25 Apr 2023 16:10:24 +0100 Subject: [PATCH 018/160] changed marker style --- indica/workflows/bayes_workflow.py | 24 ++++++++++++++++-------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/indica/workflows/bayes_workflow.py b/indica/workflows/bayes_workflow.py index 655eb5ac..449279a8 100644 --- a/indica/workflows/bayes_workflow.py +++ b/indica/workflows/bayes_workflow.py @@ -129,14 +129,22 @@ def _plot_1d( zorder=1, color="lightgrey", ) - plt.plot( - diag_data[blobkey].__getattr__(dims[0]), - diag_data[blobkey].sel(t=blob_data.t).values, - linestyle="-", - color="black", - label=f"{blobkey} data", - zorder=4, - ) + if "channel" in dims: + plt.plot( + diag_data[blobkey].__getattr__(dims[0]), + diag_data[blobkey].sel(t=blob_data.t).values, + "k^", + label=f"{blobkey} data", + zorder=4, + ) + else: + plt.plot( + diag_data[blobkey].__getattr__(dims[0]), + diag_data[blobkey].sel(t=blob_data.t).values, + "k-", + label=f"{blobkey} data", + zorder=4, + ) plt.ylabel(ylabel) plt.xlabel(dims[0]) plt.legend() From 5655cef4ffccee8413f1bea172f2ab5861e667dc Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 27 Apr 2023 09:54:57 +0100 Subject: [PATCH 019/160] small fix --- indica/workflows/bayes_workflow.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/indica/workflows/bayes_workflow.py b/indica/workflows/bayes_workflow.py index 7a4b6bab..3c0a4836 100644 --- a/indica/workflows/bayes_workflow.py +++ b/indica/workflows/bayes_workflow.py @@ -257,7 +257,7 @@ def plot_bayes_result( key, figheader=figheader, filename="temperature", - phantom_profile=phantom_profiles[key], + phantom_profile=phantom_profiles[key].sel(element="ar"), color="red", linestyle="dotted", ) @@ -270,7 +270,7 @@ def plot_bayes_result( blobs[key].sel(element="ar"), key, figheader=figheader, - phantom_profile=phantom_profiles[key], + phantom_profile=phantom_profiles[key].sel(element="ar"), color="red", ) From 40e5cb5bf7a8a5b33a0593a34edff50975224b83 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 27 Apr 2023 11:06:04 +0100 Subject: [PATCH 020/160] adjusting alpha --- indica/workflows/bayes_workflow.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/indica/workflows/bayes_workflow.py b/indica/workflows/bayes_workflow.py index 3c0a4836..fea6aae8 100644 --- a/indica/workflows/bayes_workflow.py +++ b/indica/workflows/bayes_workflow.py @@ -25,7 +25,7 @@ def plot_profile( label=f"{blobkey}, 68% Confidence", zorder=3, color=color, - alpha=0.9, + alpha=0.8, ) plt.fill_between( profile.rho_poloidal, @@ -34,7 +34,7 @@ def plot_profile( label=f"{blobkey}, 95% Confidence", zorder=2, color="grey", - alpha=0.7, + alpha=0.6, ) plt.fill_between( profile.rho_poloidal, @@ -43,7 +43,7 @@ def plot_profile( label=f"{blobkey}, Max-Min", zorder=1, color="lightgrey", - alpha=0.7, + alpha=0.6, ) if phantom_profile is not None: From 210bf7c09194eb3a029893cc3fc2a433507e3162 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Wed, 3 May 2023 11:04:07 +0100 Subject: [PATCH 021/160] black reformatting --- indica/models/equilibrium_reconstruction.py | 1 - indica/models/helike_spectroscopy.py | 111 ++++++++++---------- tests/unit/models/test_diagnostic_models.py | 1 - 3 files changed, 55 insertions(+), 58 deletions(-) diff --git a/indica/models/equilibrium_reconstruction.py b/indica/models/equilibrium_reconstruction.py index 01321da4..5d8cae2f 100644 --- a/indica/models/equilibrium_reconstruction.py +++ b/indica/models/equilibrium_reconstruction.py @@ -1,4 +1,3 @@ - import matplotlib.pyplot as plt import xarray as xr diff --git a/indica/models/helike_spectroscopy.py b/indica/models/helike_spectroscopy.py index f1af47cb..1492531a 100644 --- a/indica/models/helike_spectroscopy.py +++ b/indica/models/helike_spectroscopy.py @@ -36,17 +36,17 @@ class Helike_spectroscopy(DiagnosticModel): """ def __init__( - self, - name: str, - instrument_method="get_helike_spectroscopy", - etendue: float = 1.0, - calibration: float = 5.0e-18, - marchuk: bool = True, - full_run: bool = False, - element: str = "ar", - window_len: int = 1030, - window_lim: list = [0.394, 0.401], - window_masks: list = [], + self, + name: str, + instrument_method="get_helike_spectroscopy", + etendue: float = 1.0, + calibration: float = 5.0e-18, + marchuk: bool = True, + full_run: bool = False, + element: str = "ar", + window_len: int = 1030, + window_lim: list = [0.394, 0.401], + window_masks: list = [], ): """ Read all atomic data and initialise objects @@ -159,8 +159,8 @@ def _set_adas_pecs(self): for line in pec: pec[line]["emiss_coeff"] = ( pec[line]["emiss_coeff"] - .sel(electron_density=4.0e19, method="nearest") - .drop_vars("electron_density") + .sel(electron_density=4.0e19, method="nearest") + .drop_vars("electron_density") ) self.pec = pec @@ -231,27 +231,27 @@ def _calculate_line_emission(self, line_labels=["w", "k"]): ) for line_name in line_labels: _line_emission = ( - _pecs.isel( - line_name=(_pecs.wavelength < self.line_ranges[line_name].stop) - & (_pecs.wavelength > self.line_ranges[line_name].start) - ) - * mult + _pecs.isel( + line_name=(_pecs.wavelength < self.line_ranges[line_name].stop) + & (_pecs.wavelength > self.line_ranges[line_name].start) + ) + * mult ) _line_emission = _line_emission.sum(["type", "line_name"]) # TODO: convert PEC wavelengths to nm as per convention at TE! ev_wavelength = ph.nm_eV_conversion(nm=self.line_ranges[line_name].start) line_emission[line_name] = ( - xr.where(_line_emission >= 0, _line_emission, 0) * ev_wavelength - ) * self.calibration + xr.where(_line_emission >= 0, _line_emission, 0) * ev_wavelength + ) * self.calibration if "k" in line_emission.keys() and "w" in line_emission.keys(): line_emission["kw"] = line_emission["k"] * line_emission["w"] if "n3" in line_emission.keys() and "w" in line_emission.keys(): line_emission["n3w"] = line_emission["n345"] * line_emission["w"] if ( - "n3" in line_emission.keys() - and "n345" in line_emission.keys() - and "w" in line_emission.keys() + "n3" in line_emission.keys() + and "n345" in line_emission.keys() + and "w" in line_emission.keys() ): line_emission["tot"] = line_emission["n345"] + line_emission["w"] line_emission["n3tot"] = line_emission["n345"] * line_emission["w"] @@ -319,10 +319,10 @@ def _calculate_temperatures(self): ) def _moment_analysis( - self, - line: str, - profile_1d: DataArray = None, - half_los: bool = True, + self, + line: str, + profile_1d: DataArray = None, + half_los: bool = True, ): """ Perform moment analysis using a specific line emission as distribution function @@ -455,7 +455,7 @@ def _moment_analysis( return result, pos, err_in, err_out def _calculate_intensity( - self, + self, ): """ Returns DataArrays of emission type with co-ordinates of line label and @@ -469,8 +469,8 @@ def _calculate_intensity( _pecs_ds = _pecs.to_dataset("type") temp = [ _pecs_ds[type] - .dropna("line_name", how="all") - .interp(electron_temperature=self.Te, assume_sorted=True) + .dropna("line_name", how="all") + .interp(electron_temperature=self.Te, assume_sorted=True) for type in _pecs_ds.data_vars.keys() ] _intensity = xr.merge(temp).to_array("type") @@ -478,7 +478,7 @@ def _calculate_intensity( return intensity def _make_spectra( - self, + self, ): """ TODO: Doppler Shift / Add convolution of broadening @@ -542,7 +542,9 @@ def _build_bckc_dictionary(self): self.bckc[quant].attrs["emiss"] = self.line_emission[line] elif datatype == ("spectra", "passive"): self.bckc["spectra"] = self.measured_spectra - self.bckc["spectra"].attrs["error"] = np.sqrt(self.measured_spectra) # poisson noise + self.bckc["spectra"].attrs["error"] = np.sqrt( + self.measured_spectra + ) # poisson noise else: print(f"{quant} not available in model for {self.instrument_method}") continue @@ -562,19 +564,19 @@ def _build_bckc_dictionary(self): self.bckc["int_n3/int_tot"] = self.bckc["int_n3"] / self.bckc["int_tot"] def __call__( - self, - Te: DataArray = None, - Ti: DataArray = None, - Ne: DataArray = None, - Nimp: DataArray = None, - Fz: dict = None, - Nh: DataArray = None, - t: LabeledArray = None, - calc_spectra=True, - calc_rho: bool = False, - minimum_lines: bool = False, - moment_analysis: bool = False, - **kwargs, + self, + Te: DataArray = None, + Ti: DataArray = None, + Ne: DataArray = None, + Nimp: DataArray = None, + Fz: dict = None, + Nh: DataArray = None, + t: LabeledArray = None, + calc_spectra=True, + calc_rho: bool = False, + minimum_lines: bool = False, + moment_analysis: bool = False, + **kwargs, ): """ Calculate diagnostic measured values @@ -624,12 +626,12 @@ def __call__( Nimp = self.plasma.impurity_density.interp(t=t) else: if ( - Ne is None - or Te is None - or Nh is None - or Fz is None - or Ti is None - or Nimp is None + Ne is None + or Te is None + or Nh is None + or Fz is None + or Ti is None + or Nimp is None ): raise ValueError("Give inputs or assign plasma class!") @@ -649,10 +651,7 @@ def __call__( quant: dict = {} if moment_analysis: if minimum_lines: - line_labels = [ - "w", - "k" - ] + line_labels = ["w", "k"] names = ["int_w", "int_k", "te_kw", "ti_w"] else: line_labels = list(self.line_ranges.keys()) @@ -664,7 +663,7 @@ def __call__( "te_n3w", "te_kw", "ti_z", - "ti_w" + "ti_w", ] quant = dict(quant, **{x: self.quantities[x] for x in names}) diff --git a/tests/unit/models/test_diagnostic_models.py b/tests/unit/models/test_diagnostic_models.py index 8fa0751d..d0f6bda3 100644 --- a/tests/unit/models/test_diagnostic_models.py +++ b/tests/unit/models/test_diagnostic_models.py @@ -15,7 +15,6 @@ from indica.models.equilibrium_reconstruction import example_run as equil_recon - EXAMPLES: Dict[str, Callable] = { "bolometer_camera": bolo, "diode_filters": diodes, From efe22813d8a3faffdf060a4f9524c723c60e0245 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Wed, 3 May 2023 13:38:10 +0100 Subject: [PATCH 022/160] fixing formatting --- indica/bayesmodels.py | 8 ++++++-- tests/unit/models/test_diagnostic_models.py | 1 - 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/indica/bayesmodels.py b/indica/bayesmodels.py index e954e995..40efcc47 100644 --- a/indica/bayesmodels.py +++ b/indica/bayesmodels.py @@ -85,11 +85,15 @@ def _ln_likelihood(self): for key in self.quant_to_optimise: # Float128 since rounding of small numbers causes problems when initial results are bad fits model_data = self.bckc[key].astype("float128") - exp_data = self.data[key].sel(t=self.plasma.time_to_calculate).astype("float128") + exp_data = ( + self.data[key].sel(t=self.plasma.time_to_calculate).astype("float128") + ) exp_error = exp_data * 0.10 # Assume percentage error if none given. if hasattr(self.data[key], "error"): - if (self.data[key].error != 0).any(): # TODO: Some models have an error of 0 given + if ( + self.data[key].error != 0 + ).any(): # TODO: Some models have an error of 0 given exp_error = self.data[key].error _ln_likelihood = np.log(gaussian(model_data, exp_data, exp_error)) diff --git a/tests/unit/models/test_diagnostic_models.py b/tests/unit/models/test_diagnostic_models.py index d0f6bda3..2057e1fd 100644 --- a/tests/unit/models/test_diagnostic_models.py +++ b/tests/unit/models/test_diagnostic_models.py @@ -12,7 +12,6 @@ from indica.models.interferometry import example_run as interf from indica.models.plasma import example_run as example_plasma from indica.models.thomson_scattering import example_run as ts -from indica.models.equilibrium_reconstruction import example_run as equil_recon EXAMPLES: Dict[str, Callable] = { From 7ac95c477ba28c591e9c9c9975dd49ef78ef7df2 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Wed, 3 May 2023 13:49:29 +0100 Subject: [PATCH 023/160] flake8 fix --- indica/bayesmodels.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/indica/bayesmodels.py b/indica/bayesmodels.py index 40efcc47..0d2253f2 100644 --- a/indica/bayesmodels.py +++ b/indica/bayesmodels.py @@ -83,7 +83,8 @@ def _build_bckc(self, params: dict, **kwargs): def _ln_likelihood(self): ln_likelihood = 0 for key in self.quant_to_optimise: - # Float128 since rounding of small numbers causes problems when initial results are bad fits + # Float128 since rounding of small numbers causes problems + # when initial results are bad fits model_data = self.bckc[key].astype("float128") exp_data = ( self.data[key].sel(t=self.plasma.time_to_calculate).astype("float128") From cea61618bd9edef850044debf958ad31e9acff1d Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Wed, 3 May 2023 13:54:50 +0100 Subject: [PATCH 024/160] flake8 fix --- indica/models/helike_spectroscopy.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/indica/models/helike_spectroscopy.py b/indica/models/helike_spectroscopy.py index 1492531a..0bbefca3 100644 --- a/indica/models/helike_spectroscopy.py +++ b/indica/models/helike_spectroscopy.py @@ -280,7 +280,7 @@ def _calculate_los_integral(self, calc_rho=False): t=self.spectra.t, calc_rho=calc_rho, ) - # TODO: LOS integral removes NaNs so manually add them back (find better solution) + # TODO: LOS integral removes NaNs so add them back (find solution) self.measured_spectra[self.measured_spectra == 0] = np.nan def _calculate_temperatures(self): From 0aac384e9951e84dcba9dd56da60db2bd48f6c39 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Wed, 3 May 2023 14:30:06 +0100 Subject: [PATCH 025/160] boolean indexing -> xarray.where in call --- indica/models/helike_spectroscopy.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/indica/models/helike_spectroscopy.py b/indica/models/helike_spectroscopy.py index 6df1a992..3d1cfa92 100644 --- a/indica/models/helike_spectroscopy.py +++ b/indica/models/helike_spectroscopy.py @@ -280,9 +280,8 @@ def _calculate_los_integral(self, calc_rho=False): t=self.spectra.t, calc_rho=calc_rho, ) - # TODO: LOS integral removes NaNs so add them back (find solution) - self.measured_spectra[self.measured_spectra == 0] = np.nan - + # LOS integral removes NaNs so add them back + self.measured_spectra = self.measured_spectra.where(self.measured_spectra != 0, np.nan) def _calculate_temperatures(self): x1 = self.los_transform.x1 x1_name = self.los_transform.x1_name @@ -771,7 +770,6 @@ def example_run(pulse: int = 9229, plasma=None, plot=False, **kwargs): plasma = example_plasma( pulse=pulse, impurities=("ar",), impurity_concentration=(0.001,), n_rad=10 ) - plasma.time_to_calculate = plasma.t[5] # Create new diagnostic diagnostic_name = "xrcs" los_transform = helike_LOS_example(3) From 0259f22ede1b4c2489d6f5fbcbe2d4322f02e9b0 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Wed, 3 May 2023 15:24:05 +0100 Subject: [PATCH 026/160] fixing tests --- tests/unit/models/test_plasma.py | 30 +++++++++++++++--------------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/tests/unit/models/test_plasma.py b/tests/unit/models/test_plasma.py index 6667e367..201bcdcc 100644 --- a/tests/unit/models/test_plasma.py +++ b/tests/unit/models/test_plasma.py @@ -20,9 +20,9 @@ def setup_class(self): self.equilibrium = Equilibrium(self.equilibrium_data) self.plasma.set_equilibrium(equilibrium=self.equilibrium) - def setup_method(self): - self.plasma.electron_density = 1 # set profiles - return + # def setup_method(self): + # self.plasma.electron_density = 1 # set profiles + # return def teardown_method(self): self.plasma.initialize_variables(tstart=self.tstart, tend=self.tend, dt=self.dt) @@ -37,19 +37,19 @@ def test_plasma_volume_non_zero(self): _volume = self.plasma.volume assert len(np.where(_volume > 0)[0]) != 0 - def test_fz_is_non_zero(self): - _fz = self.plasma.fz - assert len(np.where(_fz > 0)[0]) != 0 + # def test_fz_is_non_zero(self): + # _fz = self.plasma.fz[self.impurities[0]] + # assert len(np.where(_fz > 0)[0]) != 0 + # + # def test_lz_is_non_zero(self): + # _lz_tot = self.plasma.lz_tot[self.impurities[0]] + # assert len(np.where(_lz_tot > 0)[0]) != 0 - def test_lz_is_non_zero(self): - _lz_tot = self.plasma.lz_tot - assert len(np.where(_lz_tot > 0)[0]) != 0 - - def test_fz_one_time_point(self): - return - - def test_fz_keys_match_elements(self): - return + # def test_fz_one_time_point(self): + # return + # + # def test_fz_keys_match_elements(self): + # return # class TestPlasmaProfiles: # From 8cf45671615700bafaa12e57fff9b5042b6af4a0 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Wed, 3 May 2023 15:33:10 +0100 Subject: [PATCH 027/160] removing import --- tests/integration/models/test_helike_spectroscopy.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/integration/models/test_helike_spectroscopy.py b/tests/integration/models/test_helike_spectroscopy.py index 41e58b81..670d7d80 100644 --- a/tests/integration/models/test_helike_spectroscopy.py +++ b/tests/integration/models/test_helike_spectroscopy.py @@ -1,7 +1,6 @@ import indica.models.helike_spectroscopy as helike from indica.models.helike_spectroscopy import helike_LOS_example from indica.models.plasma import example_run as example_plasma -import pytest class TestHelike: def setup_class(self): From 8170cddfd60155e4d4e99cbfe7a1327202964f1f Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Wed, 3 May 2023 15:52:56 +0100 Subject: [PATCH 028/160] precommit fixes --- indica/models/helike_spectroscopy.py | 6 +++++- .../models/test_helike_spectroscopy.py | 17 +++++++++++---- tests/integration/test_bayesmodels.py | 2 +- tests/unit/models/test_plasma.py | 21 ++++++++++++++----- 4 files changed, 35 insertions(+), 11 deletions(-) diff --git a/indica/models/helike_spectroscopy.py b/indica/models/helike_spectroscopy.py index 3d1cfa92..6741b17b 100644 --- a/indica/models/helike_spectroscopy.py +++ b/indica/models/helike_spectroscopy.py @@ -281,7 +281,10 @@ def _calculate_los_integral(self, calc_rho=False): calc_rho=calc_rho, ) # LOS integral removes NaNs so add them back - self.measured_spectra = self.measured_spectra.where(self.measured_spectra != 0, np.nan) + self.measured_spectra = self.measured_spectra.where( + self.measured_spectra != 0, np.nan + ) + def _calculate_temperatures(self): x1 = self.los_transform.x1 x1_name = self.los_transform.x1_name @@ -764,6 +767,7 @@ def helike_LOS_example(nchannels=3): ) return los_transform + def example_run(pulse: int = 9229, plasma=None, plot=False, **kwargs): # TODO: LOS sometimes crossing bad EFIT reconstruction if plasma is None: diff --git a/tests/integration/models/test_helike_spectroscopy.py b/tests/integration/models/test_helike_spectroscopy.py index 670d7d80..58440694 100644 --- a/tests/integration/models/test_helike_spectroscopy.py +++ b/tests/integration/models/test_helike_spectroscopy.py @@ -2,6 +2,7 @@ from indica.models.helike_spectroscopy import helike_LOS_example from indica.models.plasma import example_run as example_plasma + class TestHelike: def setup_class(self): self.plasma = example_plasma() # using Phantom @@ -16,22 +17,30 @@ def setup_method(self): self.model = helike.Helike_spectroscopy("diagnostic_name") self.model.set_plasma(self.plasma) - def test_helike_moment_runs_with_multiple_LOS(self, ): + def test_helike_moment_runs_with_multiple_LOS( + self, + ): self.model.set_los_transform(self.multiple_channel_los_transform) bckc = self.model(calc_spectra=False, moment_analysis=True) assert bckc - def test_helike_moment_runs_with_single_LOS(self, ): + def test_helike_moment_runs_with_single_LOS( + self, + ): self.model.set_los_transform(self.single_channel_los_transform) bckc = self.model(calc_spectra=False, moment_analysis=True) assert bckc - def test_helike_spectra_with_multiple_LOS(self, ): + def test_helike_spectra_with_multiple_LOS( + self, + ): self.model.set_los_transform(self.multiple_channel_los_transform) bckc = self.model(calc_spectra=True, moment_analysis=False) assert bckc - def test_helike_spectra_with_single_LOS(self, ): + def test_helike_spectra_with_single_LOS( + self, + ): self.model.set_los_transform(self.single_channel_los_transform) bckc = self.model(calc_spectra=True, moment_analysis=False) assert bckc diff --git a/tests/integration/test_bayesmodels.py b/tests/integration/test_bayesmodels.py index 72cb7295..b01f1aca 100644 --- a/tests/integration/test_bayesmodels.py +++ b/tests/integration/test_bayesmodels.py @@ -1,9 +1,9 @@ import emcee import flatdict -from indica.models.helike_spectroscopy import helike_LOS_example from indica.bayesmodels import BayesModels from indica.bayesmodels import get_uniform +from indica.models.helike_spectroscopy import helike_LOS_example from indica.models.helike_spectroscopy import Helike_spectroscopy from indica.models.plasma import example_run diff --git a/tests/unit/models/test_plasma.py b/tests/unit/models/test_plasma.py index 201bcdcc..9b2a8fb7 100644 --- a/tests/unit/models/test_plasma.py +++ b/tests/unit/models/test_plasma.py @@ -1,21 +1,31 @@ import numpy as np +from indica.equilibrium import Equilibrium +from indica.equilibrium import fake_equilibrium_data from indica.models.plasma import Plasma -from indica.equilibrium import fake_equilibrium_data, Equilibrium class TestPlasma: def setup_class(self): - self.tstart = 0. + self.tstart = 0.0 self.tend = 0.1 self.dt = 0.02 self.impurities = ("c", "ar", "he") self.impurity_concentration = (0.01, 0.001, 0.01) - self.plasma = Plasma(tstart=self.tstart, tend=self.tend, dt=self.dt, main_ion="h", - impurities=self.impurities, impurity_concentration=self.impurity_concentration) + self.plasma = Plasma( + tstart=self.tstart, + tend=self.tend, + dt=self.dt, + main_ion="h", + impurities=self.impurities, + impurity_concentration=self.impurity_concentration, + ) self.equilibrium_data = fake_equilibrium_data( - tstart=self.tstart, tend=self.tend, dt=self.dt, machine_dims=self.plasma.machine_dimensions + tstart=self.tstart, + tend=self.tend, + dt=self.dt, + machine_dims=self.plasma.machine_dimensions, ) self.equilibrium = Equilibrium(self.equilibrium_data) self.plasma.set_equilibrium(equilibrium=self.equilibrium) @@ -51,6 +61,7 @@ def test_plasma_volume_non_zero(self): # def test_fz_keys_match_elements(self): # return + # class TestPlasmaProfiles: # # def setup_class(self): From 96f1f5cea7b940df2a2f6007e1eeba9ef15518d2 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Fri, 12 May 2023 17:25:28 +0100 Subject: [PATCH 029/160] stashing for quick result --- indica/bayesmodels.py | 16 +- indica/models/helike_spectroscopy.py | 15 +- indica/models/plasma.py | 19 +- indica/workflows/bayes_dev_workflow.py | 334 +++++++++++++++++++++++++ indica/workflows/bayes_workflow.py | 54 +++- 5 files changed, 409 insertions(+), 29 deletions(-) create mode 100644 indica/workflows/bayes_dev_workflow.py diff --git a/indica/bayesmodels.py b/indica/bayesmodels.py index 0d2253f2..ae3666b5 100644 --- a/indica/bayesmodels.py +++ b/indica/bayesmodels.py @@ -95,10 +95,10 @@ def _ln_likelihood(self): if ( self.data[key].error != 0 ).any(): # TODO: Some models have an error of 0 given - exp_error = self.data[key].error + exp_error = self.data[key].error.sel(t=self.plasma.time_to_calculate) _ln_likelihood = np.log(gaussian(model_data, exp_data, exp_error)) - # Treat channel as key dim which isn't averaged like other dims + # treat channel as key dim which isn't averaged like other dims if "channel" in _ln_likelihood.dims: _ln_likelihood = _ln_likelihood.sum(dim="channel", skipna=True) ln_likelihood += _ln_likelihood.mean(skipna=True).values @@ -186,6 +186,8 @@ def ln_posterior(self, parameters: dict, **kwargs): return -np.inf, {} self.plasma.update_profiles(parameters) + self.plasma.calc_impurity_density("c") # Temp: put this somewhere better + self._build_bckc(parameters, **kwargs) # model calls ln_likelihood = self._ln_likelihood() # compare results to data ln_posterior = ln_likelihood + ln_prior @@ -203,6 +205,16 @@ def ln_posterior(self, parameters: dict, **kwargs): "impurity_density": self.plasma.impurity_density.sel( t=self.plasma.time_to_calculate ), + "ion_density": self.plasma.ion_density.sel( + t=self.plasma.time_to_calculate + ), + "fast_density": self.plasma.fast_density.sel( + t=self.plasma.time_to_calculate + ), + "neutral_density": self.plasma.neutral_density.sel( + t=self.plasma.time_to_calculate + ), + # TODO: add Nh } blob = deepcopy({**self.bckc, **kin_profs}) diff --git a/indica/models/helike_spectroscopy.py b/indica/models/helike_spectroscopy.py index 6741b17b..6711bee6 100644 --- a/indica/models/helike_spectroscopy.py +++ b/indica/models/helike_spectroscopy.py @@ -40,12 +40,13 @@ def __init__( name: str, instrument_method="get_helike_spectroscopy", etendue: float = 1.0, - calibration: float = 5.0e-18, + calibration: float = 5.0e-19, marchuk: bool = True, full_run: bool = False, element: str = "ar", window_len: int = 1030, window_lim: list = [0.394, 0.401], + window_vector=None, window_masks: list = [], ): """ @@ -55,8 +56,6 @@ def __init__( ---------- name String identifier for the spectrometer - fract_abu - dictionary of fractional abundance objects marchuk Use Marchuk PECs instead of ADAS adf15 files @@ -89,7 +88,11 @@ def __init__( self.pecs = self._set_adas_pecs() self.window_masks = window_masks - window = np.linspace(window_lim[0], window_lim[1], window_len) + self.window_vector = window_vector + if self.window_vector is not None: + window = self.window_vector + else: + window = np.linspace(window_lim[0], window_lim[1], window_len) mask = np.zeros(shape=window.shape) if window_masks: for mslice in window_masks: @@ -578,6 +581,7 @@ def __call__( calc_rho: bool = False, minimum_lines: bool = False, moment_analysis: bool = False, + background=None, **kwargs, ): """ @@ -682,11 +686,14 @@ def __call__( self.intensity = self._calculate_intensity() self.spectra = self._make_spectra() + self._calculate_los_integral( calc_rho=calc_rho, ) if moment_analysis: self._calculate_temperatures() + if background is not None: + self.measured_spectra = self.measured_spectra + background.sel(t=t) self._build_bckc_dictionary() return self.bckc diff --git a/indica/models/plasma.py b/indica/models/plasma.py index ddf70ce2..3bc847ef 100644 --- a/indica/models/plasma.py +++ b/indica/models/plasma.py @@ -843,22 +843,13 @@ def vloop(self): self._vloop.loc[dict(t=t)] = vloop.values return self._vloop - def calc_impurity_density(self, t=None): + def calc_impurity_density(self, elements): """ - Calculate impurity density from concentration + Calculate impurity density from concentration and electron density """ - if t is None: - t = self.t - if type(t) is not LabeledArray: - t = [t] - - profile_shape = self.Nimp_prof.yspl / self.Nimp_prof.yspl.sel(rho_poloidal=0) - for elem in self.impurities: - conc = self.impurity_concentration.sel(element=elem) - for _t in t: - dens_0 = self.electron_density.sel(rho_poloidal=0, t=t) * conc - Nimp = profile_shape * dens_0.sel(t=_t) - self.impurity_density.loc[dict(element=elem, t=_t)] = Nimp.values + for elem in elements: + Nimp = self.electron_density * self.impurity_concentration.sel(element=elem) + self.impurity_density.loc[dict(element=elem,)] = Nimp.values def impose_flat_zeff(self): """ diff --git a/indica/workflows/bayes_dev_workflow.py b/indica/workflows/bayes_dev_workflow.py new file mode 100644 index 00000000..1001cc6e --- /dev/null +++ b/indica/workflows/bayes_dev_workflow.py @@ -0,0 +1,334 @@ +import emcee +import numpy as np +import xarray as xr +import pandas as pd +import flatdict + +from indica.readers.read_st40 import ReadST40 +from indica.bayesmodels import BayesModels, get_uniform +from indica.workflows.bayes_workflow import plot_bayes_result, sample_with_autocorr + +from indica.models.interferometry import Interferometry +from indica.models.helike_spectroscopy import Helike_spectroscopy +from indica.models.charge_exchange import ChargeExchange +from indica.models.equilibrium_reconstruction import EquilibriumReconstruction +from indica.models.plasma import Plasma +from indica.converters.line_of_sight import LineOfSightTransform + +# global configurations +DEFAULT_PHANTOM_PARAMS = { + "Ne_prof.y0": 5e19, + "Ne_prof.wcenter": 0.4, + "Ne_prof.peaking": 2, + "Ne_prof.y1": 2e18, + "Ne_prof.yend": 1e18, + "Ne_prof.wped": 2, + + "Nimp_prof.y0": 3e16, + "Nimp_prof.y1": 0.5e16, + "Nimp_prof.wcenter": 0.4, + "Nimp_prof.wped": 6, + "Nimp_prof.peaking": 2, + + "Te_prof.y0": 3000, + "Te_prof.wcenter": 0.4, + "Te_prof.wped": 4, + "Te_prof.peaking": 2, + + "Ti_prof.y0": 6000, + "Ti_prof.wcenter": 0.4, + "Ti_prof.wped": 4, + "Ti_prof.peaking": 2, + } + +DEFAULT_PRIORS = { + "Ne_prof.y0": get_uniform(2e19, 4e20), + "Ne_prof.y1": get_uniform(1e18, 1e19), + "Ne_prof.y0/Ne_prof.y1": lambda x1, x2: np.where((x1 > x2 * 2), 1, 0), + "Ne_prof.wped": get_uniform(1, 6), + "Ne_prof.wcenter": get_uniform(0.1, 0.8), + "Ne_prof.peaking": get_uniform(1, 4), + + "Nimp_prof.y0": get_uniform(1e15, 1e17), + "Nimp_prof.y1": get_uniform(1e15, 2e16), + "Ne_prof.y0/Nimp_prof.y0": lambda x1, x2: np.where((x1 > x2 * 100) & (x1 < x2 * 1e5), 1, 0), + "Nimp_prof.y0/Nimp_prof.y1": lambda x1, x2: np.where((x1 > x2), 1, 0), + "Nimp_prof.wped": get_uniform(1, 6), + "Nimp_prof.wcenter": get_uniform(0.1, 0.8), + "Nimp_prof.peaking": get_uniform(1, 4), + + "Te_prof.y0": get_uniform(1000, 5000), + "Te_prof.wped": get_uniform(1, 6), + "Te_prof.wcenter": get_uniform(0.1, 0.6), + "Te_prof.peaking": get_uniform(1, 4), + "Ti_prof.y0/Te_prof.y0": lambda x1, x2: np.where(x1 > x2, 1, 0), # hot ion mode + + "Ti_prof.y0": get_uniform(1000, 10000), + "Ti_prof.wped": get_uniform(1, 6), + "Ti_prof.wcenter": get_uniform(0.1, 0.6), + "Ti_prof.peaking": get_uniform(1, 4), +} + +OPTIMISED_PARAMS = [ + "Ne_prof.y0", + # "Ne_prof.y1", + "Ne_prof.peaking", + "Ne_prof.wcenter", + "Ne_prof.wped", + + "Nimp_prof.y0", + "Nimp_prof.wcenter", + "Nimp_prof.wped", + "Nimp_prof.y1", + "Nimp_prof.peaking", + + "Te_prof.y0", + "Te_prof.wped", + "Te_prof.wcenter", + "Te_prof.peaking", + + "Ti_prof.y0", + "Ti_prof.wped", + "Ti_prof.wcenter", + "Ti_prof.peaking", +] + +OPTIMISED_QUANTITY = [ + "xrcs.spectra", + "cxff_pi.ti", + "efit.wp", + "smmh1.ne" + ] + +class BayesWorkflow: + def __init__(self, + pulse=None, + result_path="./results/example/", + nwalkers=50, + tstart=0.02, + tend=0.10, + dt=0.01, + tsample=4, + iterations=100, + burn_in=0, + diagnostics=None, + phantom=True, + profiles=None + ): + + self.pulse=pulse + self.tstart = tstart + self.tend = tend + self.dt = dt + self.tsample = tsample + self.result_path = result_path + self.nwalkers = nwalkers + self.iterations = iterations + self.burn_in = burn_in + self.diagnostics = diagnostics + self.phantom = phantom + self.profiles = profiles + + self.plasma = Plasma( + tstart=tstart, + tend=tend, + dt=dt, + main_ion="h", + impurities=("ar","c"), + impurity_concentration=(0.001, 0.04, ), + full_run=False, + n_rad=20, + ) + self.plasma.time_to_calculate = self.plasma.t[tsample] + self.plasma.update_profiles(DEFAULT_PHANTOM_PARAMS) + self.plasma.build_atomic_data(calc_power_loss=False) + + self.init_fast_particles() + self.read_st40(diagnostics) + self.init_models() + + + if self.phantom: + self.phantom_data() + self.save_profiles() + else: + self.real_data() + + self.bayes_run = BayesModels( + plasma=self.plasma, + data=self.flat_data, + diagnostic_models=[*self.models.values()], + quant_to_optimise=OPTIMISED_QUANTITY, + priors=DEFAULT_PRIORS, + ) + + ndim = len(OPTIMISED_PARAMS) + self.start_points = self.bayes_run.sample_from_priors(OPTIMISED_PARAMS, size=self.nwalkers) + self.move = [(emcee.moves.StretchMove(), 1.0), (emcee.moves.DEMove(), 0.0)] + + self.sampler = emcee.EnsembleSampler( + nwalkers, + ndim, + log_prob_fn=self.bayes_run.ln_posterior, + parameter_names=OPTIMISED_PARAMS, + moves=self.move, + kwargs={"moment_analysis": False, "calc_spectra": True, "minimum_lines": False, + "background":self.flat_data["xrcs.background"]}, + ) + + def init_fast_particles(self): + st40_code = ReadST40(13110009, self.tstart, self.tend, dt=self.dt, tree="astra") + st40_code.get_raw_data("", "astra", "RUN573") + st40_code.bin_data_in_time(["astra"], self.tstart, self.tend, self.dt) + code_data = st40_code.binned_data["astra"] + Nf = code_data["nf"].interp(rho_poloidal=self.plasma.rho, t=self.plasma.t) * 1.0e19 + self.plasma.fast_density.values = Nf.values + Nn = code_data["nn"].interp(rho_poloidal=self.plasma.rho, t=self.plasma.t) * 1.0e19 + self.plasma.neutral_density.values = Nn.values + Pblon = code_data["pblon"].interp(rho_poloidal=self.plasma.rho, t=self.plasma.t) + self.plasma.pressure_fast_parallel.values = Pblon.values + Pbper = code_data["pbper"].interp(rho_poloidal=self.plasma.rho, t=self.plasma.t) + self.plasma.pressure_fast_perpendicular.values = Pbper.values + + def read_st40(self, diagnostics=None): + self.ST40_data = ReadST40(self.pulse, tstart=self.tstart, tend=self.tend) + self.ST40_data(diagnostics) + self.plasma.set_equilibrium(self.ST40_data.equilibrium) + + def save_profiles(self): + phantom_profiles = { + "electron_density": self.plasma.electron_density.sel( + t=self.plasma.time_to_calculate + ).copy(), + "impurity_density": self.plasma.impurity_density.sel( + t=self.plasma.time_to_calculate + ).copy(), + "electron_temperature": self.plasma.electron_temperature.sel( + t=self.plasma.time_to_calculate + ).copy(), + "ion_temperature": self.plasma.ion_temperature.sel( + t=self.plasma.time_to_calculate + ).copy(), + } + self.profiles = phantom_profiles + + def init_models(self): + self.models = {} + for diag in self.diagnostics: + if diag == "smmh1": + # los_transform = self.ST40_data.binned_data["smmh1"]["ne"].transform + machine_dims = self.plasma.machine_dimensions + origin = np.array([[-0.38063365, 0.91893092, 0.01]]) + # end = np.array([[0, 0, 0.01]]) + direction = np.array([[0.38173721, -0.92387953, -0.02689453]]) + los_transform = LineOfSightTransform( + origin[:, 0], + origin[:, 1], + origin[:, 2], + direction[:, 0], + direction[:, 1], + direction[:, 2], + name="", + machine_dimensions=machine_dims, + passes=2,) + los_transform.set_equilibrium(self.plasma.equilibrium) + model = Interferometry(name="smmh1") + model.set_los_transform(los_transform) + model.plasma = self.plasma + self.models["smmh1"] = model + + if "xrcs" in self.diagnostics: + los_transform = self.ST40_data.binned_data["xrcs"]["te_kw"].transform + model = Helike_spectroscopy(name="xrcs", window_masks=[slice(0.394, 0.396)], + window_vector=self.ST40_data.binned_data["xrcs"]["spectra"].wavelength.values * 0.1) + model.set_los_transform(los_transform) + model.plasma = self.plasma + self.models["xrcs"] = model + + if "efit" in self.diagnostics: + model = EquilibriumReconstruction(name="efit") + model.plasma = self.plasma + self.models["efit"] = model + + if "cxff_pi" in self.diagnostics: + transform = self.ST40_data.binned_data["cxff_pi"]["ti"].transform + transform.set_equilibrium(self.ST40_data.equilibrium) + model = ChargeExchange(name="cxff_pi", element="ar") + model.set_transect_transform(transform) + model.plasma = self.plasma + self.models["cxff_pi"] = model + + + def phantom_data(self, noise=False, noise_factor=0.1): + + self.flat_data = {} + self.flat_data["smmh1.ne"] = self.models["smmh1"]().pop("ne").expand_dims(dim={"t": [self.plasma.time_to_calculate]}) + self.flat_data["xrcs.spectra"] = self.models["xrcs"]().pop("spectra").expand_dims(dim={"t": [self.plasma.time_to_calculate]}) + self.flat_data["xrcs.background"] = None + cxrs_data = self.models["cxff_pi"]().pop("ti").expand_dims(dim={"t": [self.plasma.time_to_calculate]}) + self.flat_data["cxff_pi.ti"] = cxrs_data.where(cxrs_data.channel == 2) + self.flat_data["efit.wp"] = self.models["efit"]().pop("wp").expand_dims(dim={"t": [self.plasma.time_to_calculate]}) + + if noise: + self.flat_data["smmh1.ne"] = self.flat_data["smmh1.ne"] + self.flat_data["smmh1.ne"].max().values * np.random.normal(0, noise_factor, None) + self.flat_data["xrcs.spectra"] = self.flat_data["xrcs.spectra"] + np.random.normal(0, np.sqrt(self.flat_data["xrcs.spectra"].values[0,]), self.flat_data["xrcs.spectra"].shape[1]) + self.flat_data["cxff_pi.ti"] = self.flat_data["cxff_pi.ti"] + self.flat_data["cxff_pi.ti"].max().values * np.random.normal(0, noise_factor, self.flat_data["cxff_pi.ti"].shape[1]) + self.flat_data["efit.wp"] = self.flat_data["efit.wp"] + self.flat_data["efit.wp"].max().values * np.random.normal(0, noise_factor, None) + + + def real_data(self): + + self.flat_data = flatdict.FlatDict(self.ST40_data.binned_data, ".") + if "xrcs" in self.diagnostics: + self.flat_data["xrcs.spectra"]["wavelength"] = self.flat_data["xrcs.spectra"].wavelength * 0.1 + background = self.flat_data["xrcs.spectra"].where( + (self.flat_data["xrcs.spectra"].wavelength < 0.392) & + (self.flat_data["xrcs.spectra"].wavelength > 0.388), + drop=True) + self.flat_data["xrcs.background"] = background.mean(dim="wavelength") + self.flat_data["xrcs.spectra"]["error"] = np.sqrt(self.flat_data["xrcs.spectra"] + background.std(dim="wavelength") ** 2) + + if "cxff_pi" in self.diagnostics: + self.flat_data["cxff_pi"]["ti"] = self.flat_data["cxff_pi"]["ti"].where(self.flat_data["cxff_pi"]["ti"].channel==2, drop=True) + + + def __call__(self, *args, **kwargs): + + autocorr = sample_with_autocorr( + self.sampler, self.start_points, iterations=self.iterations, auto_sample=5 + ) + blobs = self.sampler.get_blobs(discard=self.burn_in, flat=True) + blob_names = self.sampler.get_blobs().flatten()[0].keys() + blob_dict = { + blob_name: xr.concat( + [data[blob_name] for data in blobs], + dim=pd.Index(np.arange(0, blobs.__len__()), name="index"), + ) + for blob_name in blob_names + } + samples = self.sampler.get_chain(flat=True) + + prior_samples = self.bayes_run.sample_from_priors(OPTIMISED_PARAMS, size=int(1e5)) + result = { + "blobs": blob_dict, + "diag_data": self.flat_data, + "samples": samples, + "prior_samples": prior_samples, + "param_names": OPTIMISED_PARAMS, + "phantom_profiles": self.profiles, + "autocorr": autocorr, + } + self.acceptance_fraction = self.sampler.acceptance_fraction.sum() + print(self.acceptance_fraction) + plot_bayes_result(**result, figheader=self.result_path) + +if __name__ == "__main__": + + run = BayesWorkflow(pulse=10009, result_path="./results/10009_full/", iterations=10000, nwalkers=200, + burn_in=100, diagnostics=[ + "xrcs", + "efit", + "smmh1", + "cxff_pi" + ], phantom=False) + run() diff --git a/indica/workflows/bayes_workflow.py b/indica/workflows/bayes_workflow.py index fea6aae8..b12e1986 100644 --- a/indica/workflows/bayes_workflow.py +++ b/indica/workflows/bayes_workflow.py @@ -47,7 +47,11 @@ def plot_profile( ) if phantom_profile is not None: - phantom_profile.plot( + if "element" in phantom_profile[blobkey].dims: + phantom = phantom_profile[blobkey].sel(element="ar") + else: + phantom = phantom_profile[blobkey] + phantom.plot( label=f"{blobkey}, phantom profile", linestyle=linestyle, color="black", @@ -73,11 +77,12 @@ def _plot_0d( figheader="./results/test/", xlabel="samples ()", ylabel="a.u.", + figsize=(6.4, 4.8), **kwargs, ): if blobkey not in blobs.keys(): raise ValueError(f"{blobkey} not in blobs") - plt.figure() + plt.figure(figsize=figsize) blob_data = blobs[blobkey] plt.plot(blob_data, label=f"{blobkey} model") plt.axhline( @@ -100,12 +105,14 @@ def _plot_1d( filename: str, figheader="./results/test/", ylabel="a.u.", + xlim = None, + figsize=(6.4, 4.8), **kwargs, ): if blobkey not in blobs.keys(): raise ValueError(f"{blobkey} not in blobs") - plt.figure() + plt.figure(figsize=figsize) blob_data = blobs[blobkey] dims = tuple(name for name in blob_data.dims if name != "index") plt.fill_between( @@ -150,6 +157,7 @@ def _plot_1d( ) plt.ylabel(ylabel) plt.xlabel(dims[0]) + plt.xlim(xlim) plt.legend() plt.savefig(figheader + filename) plt.close() @@ -229,8 +237,10 @@ def plot_bayes_result( f"{key.replace('.', '_')}.png", figheader=figheader, ylabel="intensity (A.U.)", + xlim = (0.394, 0.401), + figsize=(12, 10), ) - key = "cxrs.ti" + key = "cxff_pi.ti" if key in blobs.keys(): _plot_1d( blobs, @@ -246,7 +256,7 @@ def plot_bayes_result( blobs[key], key, figheader=figheader, - phantom_profile=phantom_profiles[key], + phantom_profile=phantom_profiles, sharefig=True, color="blue", linestyle="dashdot", @@ -257,22 +267,48 @@ def plot_bayes_result( key, figheader=figheader, filename="temperature", - phantom_profile=phantom_profiles[key].sel(element="ar"), + phantom_profile=phantom_profiles, color="red", linestyle="dotted", ) key = "electron_density" plot_profile( - blobs[key], key, figheader=figheader, phantom_profile=phantom_profiles[key] + blobs[key], key, figheader=figheader, phantom_profile=phantom_profiles ) key = "impurity_density" + for elem in blobs[key].element.values: + plot_profile( + blobs[key].sel(element=elem), + key, + figheader=figheader, + filename=f"{elem} density", + phantom_profile=phantom_profiles, + color="red", + ) + key = "ion_density" plot_profile( - blobs[key].sel(element="ar"), + blobs[key].sel(element="h"), key, figheader=figheader, - phantom_profile=phantom_profiles[key].sel(element="ar"), + filename=f"h density", + phantom_profile=phantom_profiles, color="red", ) + key = "fast_density" + plot_profile( + blobs[key], + key, + figheader=figheader, + phantom_profile=phantom_profiles, + color="red", + ) + key = "neutral_density" + plot_profile( + blobs[key], + key, + figheader=figheader, + phantom_profile=phantom_profiles, + ) corner.corner(samples, labels=param_names) plt.savefig(figheader + "posterior.png") From 7e1fda10f78fd12d0530c417919f440eb6173870 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Mon, 22 May 2023 10:40:32 +0100 Subject: [PATCH 030/160] added calc_impurity for generating profiles from concentrations to update profiles method --- indica/models/plasma.py | 162 ++++++++++++++++++++-------------------- 1 file changed, 82 insertions(+), 80 deletions(-) diff --git a/indica/models/plasma.py b/indica/models/plasma.py index 3bc847ef..83373f7b 100644 --- a/indica/models/plasma.py +++ b/indica/models/plasma.py @@ -103,18 +103,18 @@ class Plasma: def __init__( - self, - tstart: float = 0.01, - tend: float = 0.14, - dt: float = 0.01, - machine_dimensions=((0.15, 0.95), (-0.7, 0.7)), - impurities: tuple = ("c", "ar"), - main_ion: str = "h", - impurity_concentration: tuple = (0.02, 0.001), - pulse: int = None, - full_run: bool = False, - n_rad: int = 41, - verbose: bool = False, + self, + tstart: float = 0.01, + tend: float = 0.14, + dt: float = 0.01, + machine_dimensions=((0.15, 0.95), (-0.7, 0.7)), + impurities: tuple = ("c", "ar"), + main_ion: str = "h", + impurity_concentration: tuple = (0.02, 0.001), + pulse: int = None, + full_run: bool = False, + n_rad: int = 41, + verbose: bool = False, ): """ Class for plasma objects. @@ -458,10 +458,12 @@ def initialize_variables(self, tstart: float, tend: float, dt: float): ) def assign_profiles( - self, profile: str = "electron_density", t: float = None, element: str = "ar" + self, profile: str = "electron_density", t: float = None, element: str = "ar" ): if profile == "electron_density": self.electron_density.loc[dict(t=t)] = self.Ne_prof() + self.calc_impurity_density( + tuple([imp for imp in self.impurities if imp != element])) # only update non-argon elif profile == "electron_temperature": self.electron_temperature.loc[dict(t=t)] = self.Te_prof() elif profile == "ion_temperature": @@ -478,15 +480,15 @@ def assign_profiles( ) def update_profiles( - self, - parameters: dict, - profile_prefixs: list = [ - "Te_prof", - "Ti_prof", - "Ne_prof", - "Nimp_prof", - "Vrot_prof", - ], + self, + parameters: dict, + profile_prefixs: list = [ + "Te_prof", + "Ti_prof", + "Ne_prof", + "Nimp_prof", + "Vrot_prof", + ], ): """ Update plasma profiles with profile parameters i.e. @@ -550,7 +552,7 @@ def pressure_tot(self): def pressure_fast(self): # TODO: check whether degrees of freedom are correctly included... self._pressure_fast.values = ( - self.pressure_fast_parallel / 3 + self.pressure_fast_perpendicular * 2 / 3 + self.pressure_fast_parallel / 3 + self.pressure_fast_perpendicular * 2 / 3 ) return self._pressure_fast @@ -617,8 +619,8 @@ def calc_zeff(self): meanz = self.meanz for elem in self.elements: self._zeff.loc[dict(element=elem)] = ( - (ion_density.sel(element=elem) * meanz.sel(element=elem) ** 2) - / electron_density + (ion_density.sel(element=elem) * meanz.sel(element=elem) ** 2) + / electron_density ).values return self._zeff @@ -662,8 +664,8 @@ def calc_lz_tot(self): self.power_loss_tot[elem]( Te, Fz, Ne=Ne, Nh=Nh, full_run=self.full_run ) - .transpose() - .values + .transpose() + .values ) return self._lz_tot @@ -690,8 +692,8 @@ def calc_lz_sxr(self): self.power_loss_sxr[elem]( Te, Fz, Ne=Ne, Nh=Nh, full_run=self.full_run ) - .transpose() - .values + .transpose() + .values ) return self._lz_sxr @@ -704,9 +706,9 @@ def calc_total_radiation(self): ion_density = self.ion_density for elem in self.elements: total_radiation = ( - lz_tot[elem].sum("ion_charges") - * self.electron_density - * ion_density.sel(element=elem) + lz_tot[elem].sum("ion_charges") + * self.electron_density + * ion_density.sel(element=elem) ) self._total_radiation.loc[dict(element=elem)] = xr.where( total_radiation >= 0, @@ -727,9 +729,9 @@ def calc_sxr_radiation(self): ion_density = self.ion_density for elem in self.elements: sxr_radiation = ( - lz_sxr[elem].sum("ion_charges") - * self.electron_density - * ion_density.sel(element=elem) + lz_sxr[elem].sum("ion_charges") + * self.electron_density + * ion_density.sel(element=elem) ) self._sxr_radiation.loc[dict(element=elem)] = xr.where( sxr_radiation >= 0, @@ -849,7 +851,7 @@ def calc_impurity_density(self, elements): """ for elem in elements: Nimp = self.electron_density * self.impurity_concentration.sel(element=elem) - self.impurity_density.loc[dict(element=elem,)] = Nimp.values + self.impurity_density.loc[dict(element=elem, )] = Nimp.values def impose_flat_zeff(self): """ @@ -859,14 +861,14 @@ def impose_flat_zeff(self): for elem in self.impurities: if np.count_nonzero(self.ion_density.sel(element=elem)) != 0: zeff_tmp = ( - self.ion_density.sel(element=elem) - * self.meanz.sel(element=elem) ** 2 - / self.electron_density + self.ion_density.sel(element=elem) + * self.meanz.sel(element=elem) ** 2 + / self.electron_density ) value = zeff_tmp.where(zeff_tmp.rho_poloidal < 0.2).mean("rho_poloidal") zeff_tmp = zeff_tmp / zeff_tmp * value ion_density_tmp = zeff_tmp / ( - self.meanz.sel(element=elem) ** 2 / self.electron_density + self.meanz.sel(element=elem) ** 2 / self.electron_density ) self.ion_density.loc[dict(element=elem)] = ion_density_tmp.values @@ -878,13 +880,13 @@ def convert_in_time(self, value: DataArray, method="linear"): return binned def build_atomic_data( - self, - Te: DataArray = None, - Ne: DataArray = None, - Nh: DataArray = None, - tau: DataArray = None, - default=True, - calc_power_loss=True, + self, + Te: DataArray = None, + Ne: DataArray = None, + Nh: DataArray = None, + tau: DataArray = None, + default=True, + calc_power_loss=True, ): if default: xend = 1.02 @@ -928,8 +930,8 @@ def build_atomic_data( power_loss_tot[elem] = PowerLoss(plt, prb, PRC=prc) power_loss_tot[elem](Te, F_z_t, Ne=Ne, Nh=Nh, full_run=self.full_run) if ( - "pls" in self.adf11[elem].keys() - and "prs" in self.adf11[elem].keys() + "pls" in self.adf11[elem].keys() + and "prs" in self.adf11[elem].keys() ): try: pls = self.ADASReader.get_adf11( @@ -995,13 +997,13 @@ def map_to_midplane(self): for t in np.array(self.time_to_calculate, ndmin=1): rho = ( self.equilibrium.rho.sel(t=t, method="nearest") - .interp(R=R, z=z) - .drop_vars(["R", "z"]) + .interp(R=R, z=z) + .drop_vars(["R", "z"]) ) midplane_profiles[k].append( prof_rho.sel(t=t, method="nearest") - .interp(rho_poloidal=rho) - .drop_vars("rho_poloidal") + .interp(rho_poloidal=rho) + .drop_vars("rho_poloidal") ) midplane_profiles[k] = xr.concat(midplane_profiles[k], "t").assign_coords( t=self.t @@ -1013,7 +1015,7 @@ def map_to_midplane(self): self.midplane_profiles = midplane_profiles def calc_centrifugal_asymmetry( - self, time=None, test_toroidal_rotation=None, plot=False + self, time=None, test_toroidal_rotation=None, plot=False ): """ Calculate (R, z) maps of the ion densities caused by centrifugal asymmetry @@ -1072,12 +1074,12 @@ def calc_centrifugal_asymmetry( self.centrifugal_asymmetry.loc[dict(element=elem)] = asymm asymmetry_factor = asymm.interp(rho_poloidal=self.rho_2d) self.asymmetry_multiplier.loc[dict(element=elem)] = np.exp( - asymmetry_factor * (self.rho_2d.R**2 - R_0**2) + asymmetry_factor * (self.rho_2d.R ** 2 - R_0 ** 2) ) self.ion_density_2d = ( - ion_density.interp(rho_poloidal=self.rho_2d).drop_vars("rho_poloidal") - * self.asymmetry_multiplier + ion_density.interp(rho_poloidal=self.rho_2d).drop_vars("rho_poloidal") + * self.asymmetry_multiplier ) assign_datatype(self.ion_density_2d, ("density", "ion"), "m^-3") @@ -1118,9 +1120,9 @@ def calc_rad_power_2d(self): """ for elem in self.elements: total_radiation = ( - self.lz_tot[elem].sum("ion_charges") - * self.electron_density - * self.ion_density.sel(element=elem) + self.lz_tot[elem].sum("ion_charges") + * self.electron_density + * self.ion_density.sel(element=elem) ) total_radiation = xr.where( total_radiation >= 0, @@ -1130,9 +1132,9 @@ def calc_rad_power_2d(self): self.total_radiation.loc[dict(element=elem)] = total_radiation.values sxr_radiation = ( - self.lz_sxr[elem].sum("ion_charges") - * self.electron_density - * self.ion_density.sel(element=elem) + self.lz_sxr[elem].sum("ion_charges") + * self.electron_density + * self.ion_density.sel(element=elem) ) sxr_radiation = xr.where( sxr_radiation >= 0, @@ -1169,9 +1171,9 @@ def write_to_pickle(self): # Generalized dependency caching class TrackDependecies: def __init__( - self, - operator: Callable, - dependencies: list, + self, + operator: Callable, + dependencies: list, ): """ Call operator only if dependencies variables have changed. @@ -1188,8 +1190,8 @@ def __init__( self.dependencies = dependencies def numpyhash( - self, - nparray: np.array, + self, + nparray: np.array, ): a = nparray.view(np.uint8) return hashlib.sha1(a).hexdigest() @@ -1231,18 +1233,18 @@ def __call__(self): def example_run( - pulse: int = None, - tstart=0.02, - tend=0.1, - dt=0.01, - main_ion="h", - impurities=("c", "ar", "he"), - impurity_concentration=(0.03, 0.001, 0.01), - verbose: bool = True, - n_rad: int = 41, - full_run=False, - calc_power_loss: bool = False, - **kwargs, + pulse: int = None, + tstart=0.02, + tend=0.1, + dt=0.01, + main_ion="h", + impurities=("c", "ar", "he"), + impurity_concentration=(0.03, 0.001, 0.01), + verbose: bool = True, + n_rad: int = 41, + full_run=False, + calc_power_loss: bool = False, + **kwargs, ): # TODO: swap all profiles to new version! From 402cd2d26416ac5cd32e8a108e77909d188b3364 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Mon, 22 May 2023 10:58:23 +0100 Subject: [PATCH 031/160] added conditional handling of diagnostics in data pipeline --- indica/workflows/bayes_dev_workflow.py | 50 +++++++++++++++----------- 1 file changed, 29 insertions(+), 21 deletions(-) diff --git a/indica/workflows/bayes_dev_workflow.py b/indica/workflows/bayes_dev_workflow.py index 1001cc6e..114ed166 100644 --- a/indica/workflows/bayes_dev_workflow.py +++ b/indica/workflows/bayes_dev_workflow.py @@ -47,7 +47,7 @@ "Ne_prof.y0/Ne_prof.y1": lambda x1, x2: np.where((x1 > x2 * 2), 1, 0), "Ne_prof.wped": get_uniform(1, 6), "Ne_prof.wcenter": get_uniform(0.1, 0.8), - "Ne_prof.peaking": get_uniform(1, 4), + "Ne_prof.peaking": get_uniform(1, 6), "Nimp_prof.y0": get_uniform(1e15, 1e17), "Nimp_prof.y1": get_uniform(1e15, 2e16), @@ -55,23 +55,23 @@ "Nimp_prof.y0/Nimp_prof.y1": lambda x1, x2: np.where((x1 > x2), 1, 0), "Nimp_prof.wped": get_uniform(1, 6), "Nimp_prof.wcenter": get_uniform(0.1, 0.8), - "Nimp_prof.peaking": get_uniform(1, 4), + "Nimp_prof.peaking": get_uniform(1, 6), "Te_prof.y0": get_uniform(1000, 5000), "Te_prof.wped": get_uniform(1, 6), "Te_prof.wcenter": get_uniform(0.1, 0.6), - "Te_prof.peaking": get_uniform(1, 4), + "Te_prof.peaking": get_uniform(1, 6), "Ti_prof.y0/Te_prof.y0": lambda x1, x2: np.where(x1 > x2, 1, 0), # hot ion mode - "Ti_prof.y0": get_uniform(1000, 10000), + "Ti_prof.y0": get_uniform(2000, 10000), "Ti_prof.wped": get_uniform(1, 6), "Ti_prof.wcenter": get_uniform(0.1, 0.6), - "Ti_prof.peaking": get_uniform(1, 4), + "Ti_prof.peaking": get_uniform(1, 6), } OPTIMISED_PARAMS = [ "Ne_prof.y0", - # "Ne_prof.y1", + "Ne_prof.y1", "Ne_prof.peaking", "Ne_prof.wcenter", "Ne_prof.wped", @@ -116,7 +116,7 @@ def __init__(self, profiles=None ): - self.pulse=pulse + self.pulse = pulse self.tstart = tstart self.tend = tend self.dt = dt @@ -134,7 +134,7 @@ def __init__(self, tend=tend, dt=dt, main_ion="h", - impurities=("ar","c"), + impurities=("ar", "c"), impurity_concentration=(0.001, 0.04, ), full_run=False, n_rad=20, @@ -152,7 +152,7 @@ def __init__(self, self.phantom_data() self.save_profiles() else: - self.real_data() + self.exp_data() self.bayes_run = BayesModels( plasma=self.plasma, @@ -164,7 +164,7 @@ def __init__(self, ndim = len(OPTIMISED_PARAMS) self.start_points = self.bayes_run.sample_from_priors(OPTIMISED_PARAMS, size=self.nwalkers) - self.move = [(emcee.moves.StretchMove(), 1.0), (emcee.moves.DEMove(), 0.0)] + self.move = [(emcee.moves.StretchMove(), 0.9), (emcee.moves.DEMove(), 0.1)] self.sampler = emcee.EnsembleSampler( nwalkers, @@ -209,6 +209,12 @@ def save_profiles(self): "ion_temperature": self.plasma.ion_temperature.sel( t=self.plasma.time_to_calculate ).copy(), + "fast_density": self.plasma.fast_density.sel( + t=self.plasma.time_to_calculate + ).copy(), + "neutral_density": self.plasma.neutral_density.sel( + t=self.plasma.time_to_calculate + ).copy(), } self.profiles = phantom_profiles @@ -260,14 +266,17 @@ def init_models(self): def phantom_data(self, noise=False, noise_factor=0.1): - self.flat_data = {} - self.flat_data["smmh1.ne"] = self.models["smmh1"]().pop("ne").expand_dims(dim={"t": [self.plasma.time_to_calculate]}) - self.flat_data["xrcs.spectra"] = self.models["xrcs"]().pop("spectra").expand_dims(dim={"t": [self.plasma.time_to_calculate]}) - self.flat_data["xrcs.background"] = None - cxrs_data = self.models["cxff_pi"]().pop("ti").expand_dims(dim={"t": [self.plasma.time_to_calculate]}) - self.flat_data["cxff_pi.ti"] = cxrs_data.where(cxrs_data.channel == 2) - self.flat_data["efit.wp"] = self.models["efit"]().pop("wp").expand_dims(dim={"t": [self.plasma.time_to_calculate]}) + if "smmh1" in self.diagnostics: + self.flat_data["smmh1.ne"] = self.models["smmh1"]().pop("ne").expand_dims(dim={"t": [self.plasma.time_to_calculate]}) + if "xrcs" in self.diagnostics: + self.flat_data["xrcs.spectra"] = self.models["xrcs"]().pop("spectra").expand_dims(dim={"t": [self.plasma.time_to_calculate]}) + self.flat_data["xrcs.background"] = None + if "cxff_pi" in self.diagnostics: + cxrs_data = self.models["cxff_pi"]().pop("ti").expand_dims(dim={"t": [self.plasma.time_to_calculate]}) + self.flat_data["cxff_pi.ti"] = cxrs_data.where(cxrs_data.channel == 2) + if "efit" in self.diagnostics: + self.flat_data["efit.wp"] = self.models["efit"]().pop("wp").expand_dims(dim={"t": [self.plasma.time_to_calculate]}) if noise: self.flat_data["smmh1.ne"] = self.flat_data["smmh1.ne"] + self.flat_data["smmh1.ne"].max().values * np.random.normal(0, noise_factor, None) @@ -276,8 +285,7 @@ def phantom_data(self, noise=False, noise_factor=0.1): self.flat_data["efit.wp"] = self.flat_data["efit.wp"] + self.flat_data["efit.wp"].max().values * np.random.normal(0, noise_factor, None) - def real_data(self): - + def exp_data(self): self.flat_data = flatdict.FlatDict(self.ST40_data.binned_data, ".") if "xrcs" in self.diagnostics: self.flat_data["xrcs.spectra"]["wavelength"] = self.flat_data["xrcs.spectra"].wavelength * 0.1 @@ -324,8 +332,8 @@ def __call__(self, *args, **kwargs): if __name__ == "__main__": - run = BayesWorkflow(pulse=10009, result_path="./results/10009_full/", iterations=10000, nwalkers=200, - burn_in=100, diagnostics=[ + run = BayesWorkflow(pulse=10009, result_path="./results/10009_test/", iterations=1000, nwalkers=200, + burn_in=50, diagnostics=[ "xrcs", "efit", "smmh1", From fb7704fe846392de3221011e215f4fd6980ffd71 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Mon, 22 May 2023 10:58:59 +0100 Subject: [PATCH 032/160] moved calc_impurity to plasma method --- indica/bayesmodels.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/indica/bayesmodels.py b/indica/bayesmodels.py index ae3666b5..7d56c63e 100644 --- a/indica/bayesmodels.py +++ b/indica/bayesmodels.py @@ -186,8 +186,6 @@ def ln_posterior(self, parameters: dict, **kwargs): return -np.inf, {} self.plasma.update_profiles(parameters) - self.plasma.calc_impurity_density("c") # Temp: put this somewhere better - self._build_bckc(parameters, **kwargs) # model calls ln_likelihood = self._ln_likelihood() # compare results to data ln_posterior = ln_likelihood + ln_prior From 4997194068ca6bf584caa00abe54226503fed964 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Mon, 22 May 2023 15:25:59 +0100 Subject: [PATCH 033/160] tweaking calibration factor --- indica/models/helike_spectroscopy.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/indica/models/helike_spectroscopy.py b/indica/models/helike_spectroscopy.py index 6711bee6..74404e1b 100644 --- a/indica/models/helike_spectroscopy.py +++ b/indica/models/helike_spectroscopy.py @@ -40,7 +40,7 @@ def __init__( name: str, instrument_method="get_helike_spectroscopy", etendue: float = 1.0, - calibration: float = 5.0e-19, + calibration: float = 5.0e-20, marchuk: bool = True, full_run: bool = False, element: str = "ar", From c8e821f03ec2a032ec3e9f58ee82f53160d3a65e Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Mon, 22 May 2023 16:03:16 +0100 Subject: [PATCH 034/160] tweaking calibration factor --- indica/models/helike_spectroscopy.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/indica/models/helike_spectroscopy.py b/indica/models/helike_spectroscopy.py index 74404e1b..0c950965 100644 --- a/indica/models/helike_spectroscopy.py +++ b/indica/models/helike_spectroscopy.py @@ -40,7 +40,7 @@ def __init__( name: str, instrument_method="get_helike_spectroscopy", etendue: float = 1.0, - calibration: float = 5.0e-20, + calibration: float = 2.0e-20, marchuk: bool = True, full_run: bool = False, element: str = "ar", From 5bb0b516f5e352c95bc10eb166c5b3685d5053e5 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Mon, 22 May 2023 17:02:55 +0100 Subject: [PATCH 035/160] tweaking calibration factor --- indica/models/helike_spectroscopy.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/indica/models/helike_spectroscopy.py b/indica/models/helike_spectroscopy.py index 0c950965..eae06066 100644 --- a/indica/models/helike_spectroscopy.py +++ b/indica/models/helike_spectroscopy.py @@ -40,7 +40,7 @@ def __init__( name: str, instrument_method="get_helike_spectroscopy", etendue: float = 1.0, - calibration: float = 2.0e-20, + calibration: float = 4.0e-20, marchuk: bool = True, full_run: bool = False, element: str = "ar", From 457da6b5e63e728c96ca2cb02ea2267a1283c7da Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 23 May 2023 10:06:08 +0100 Subject: [PATCH 036/160] changed burn_in to burn_in_fraction --- indica/workflows/bayes_dev_workflow.py | 163 +++++++++++++------------ 1 file changed, 88 insertions(+), 75 deletions(-) diff --git a/indica/workflows/bayes_dev_workflow.py b/indica/workflows/bayes_dev_workflow.py index 114ed166..b7a094a0 100644 --- a/indica/workflows/bayes_dev_workflow.py +++ b/indica/workflows/bayes_dev_workflow.py @@ -4,6 +4,8 @@ import pandas as pd import flatdict +from scipy.stats import loguniform + from indica.readers.read_st40 import ReadST40 from indica.bayesmodels import BayesModels, get_uniform from indica.workflows.bayes_workflow import plot_bayes_result, sample_with_autocorr @@ -17,29 +19,29 @@ # global configurations DEFAULT_PHANTOM_PARAMS = { - "Ne_prof.y0": 5e19, - "Ne_prof.wcenter": 0.4, - "Ne_prof.peaking": 2, - "Ne_prof.y1": 2e18, - "Ne_prof.yend": 1e18, - "Ne_prof.wped": 2, - - "Nimp_prof.y0": 3e16, - "Nimp_prof.y1": 0.5e16, - "Nimp_prof.wcenter": 0.4, - "Nimp_prof.wped": 6, - "Nimp_prof.peaking": 2, - - "Te_prof.y0": 3000, - "Te_prof.wcenter": 0.4, - "Te_prof.wped": 4, - "Te_prof.peaking": 2, - - "Ti_prof.y0": 6000, - "Ti_prof.wcenter": 0.4, - "Ti_prof.wped": 4, - "Ti_prof.peaking": 2, - } + "Ne_prof.y0": 5e19, + "Ne_prof.wcenter": 0.4, + "Ne_prof.peaking": 2, + "Ne_prof.y1": 2e18, + "Ne_prof.yend": 1e18, + "Ne_prof.wped": 2, + + "Nimp_prof.y0": 3e16, + "Nimp_prof.y1": 0.5e16, + "Nimp_prof.wcenter": 0.4, + "Nimp_prof.wped": 6, + "Nimp_prof.peaking": 2, + + "Te_prof.y0": 3000, + "Te_prof.wcenter": 0.4, + "Te_prof.wped": 4, + "Te_prof.peaking": 2, + + "Ti_prof.y0": 6000, + "Ti_prof.wcenter": 0.4, + "Ti_prof.wped": 4, + "Ti_prof.peaking": 2, +} DEFAULT_PRIORS = { "Ne_prof.y0": get_uniform(2e19, 4e20), @@ -49,6 +51,8 @@ "Ne_prof.wcenter": get_uniform(0.1, 0.8), "Ne_prof.peaking": get_uniform(1, 6), + "ar_conc": loguniform(0.0001, 0.01), + "Nimp_prof.y0": get_uniform(1e15, 1e17), "Nimp_prof.y1": get_uniform(1e15, 2e16), "Ne_prof.y0/Nimp_prof.y0": lambda x1, x2: np.where((x1 > x2 * 100) & (x1 < x2 * 1e5), 1, 0), @@ -56,6 +60,7 @@ "Nimp_prof.wped": get_uniform(1, 6), "Nimp_prof.wcenter": get_uniform(0.1, 0.8), "Nimp_prof.peaking": get_uniform(1, 6), + "Nimp_prof.peaking/Ne_prof.peaking": lambda x1, x2: np.where((x1 > x2), 1, 0), # impurity always more peaked "Te_prof.y0": get_uniform(1000, 5000), "Te_prof.wped": get_uniform(1, 6), @@ -71,34 +76,36 @@ OPTIMISED_PARAMS = [ "Ne_prof.y0", - "Ne_prof.y1", + # "Ne_prof.y1", "Ne_prof.peaking", "Ne_prof.wcenter", "Ne_prof.wped", + # "ar_conc", "Nimp_prof.y0", "Nimp_prof.wcenter", "Nimp_prof.wped", - "Nimp_prof.y1", + # "Nimp_prof.y1", "Nimp_prof.peaking", "Te_prof.y0", - "Te_prof.wped", - "Te_prof.wcenter", - "Te_prof.peaking", + # "Te_prof.wped", + # "Te_prof.wcenter", + # "Te_prof.peaking", "Ti_prof.y0", - "Ti_prof.wped", - "Ti_prof.wcenter", - "Ti_prof.peaking", + # "Ti_prof.wped", + # "Ti_prof.wcenter", + # "Ti_prof.peaking", ] OPTIMISED_QUANTITY = [ - "xrcs.spectra", - "cxff_pi.ti", - "efit.wp", - "smmh1.ne" - ] + "xrcs.spectra", + "cxff_pi.ti", + "efit.wp", + "smmh1.ne" +] + class BayesWorkflow: def __init__(self, @@ -110,7 +117,7 @@ def __init__(self, dt=0.01, tsample=4, iterations=100, - burn_in=0, + burn_in_fraction=0, diagnostics=None, phantom=True, profiles=None @@ -124,7 +131,7 @@ def __init__(self, self.result_path = result_path self.nwalkers = nwalkers self.iterations = iterations - self.burn_in = burn_in + self.burn_in_fraction = burn_in_fraction self.diagnostics = diagnostics self.phantom = phantom self.profiles = profiles @@ -135,7 +142,7 @@ def __init__(self, dt=dt, main_ion="h", impurities=("ar", "c"), - impurity_concentration=(0.001, 0.04, ), + impurity_concentration=(0.001, 0.04,), full_run=False, n_rad=20, ) @@ -147,7 +154,6 @@ def __init__(self, self.read_st40(diagnostics) self.init_models() - if self.phantom: self.phantom_data() self.save_profiles() @@ -155,12 +161,12 @@ def __init__(self, self.exp_data() self.bayes_run = BayesModels( - plasma=self.plasma, - data=self.flat_data, - diagnostic_models=[*self.models.values()], - quant_to_optimise=OPTIMISED_QUANTITY, - priors=DEFAULT_PRIORS, - ) + plasma=self.plasma, + data=self.flat_data, + diagnostic_models=[*self.models.values()], + quant_to_optimise=OPTIMISED_QUANTITY, + priors=DEFAULT_PRIORS, + ) ndim = len(OPTIMISED_PARAMS) self.start_points = self.bayes_run.sample_from_priors(OPTIMISED_PARAMS, size=self.nwalkers) @@ -173,7 +179,7 @@ def __init__(self, parameter_names=OPTIMISED_PARAMS, moves=self.move, kwargs={"moment_analysis": False, "calc_spectra": True, "minimum_lines": False, - "background":self.flat_data["xrcs.background"]}, + "background": self.flat_data["xrcs.background"]}, ) def init_fast_particles(self): @@ -224,9 +230,9 @@ def init_models(self): if diag == "smmh1": # los_transform = self.ST40_data.binned_data["smmh1"]["ne"].transform machine_dims = self.plasma.machine_dimensions - origin = np.array([[-0.38063365, 0.91893092, 0.01]]) + origin = np.array([[-0.38063365, 0.91893092, 0.01]]) # end = np.array([[0, 0, 0.01]]) - direction = np.array([[0.38173721, -0.92387953, -0.02689453]]) + direction = np.array([[0.38173721, -0.92387953, -0.02689453]]) los_transform = LineOfSightTransform( origin[:, 0], origin[:, 1], @@ -236,7 +242,7 @@ def init_models(self): direction[:, 2], name="", machine_dimensions=machine_dims, - passes=2,) + passes=2, ) los_transform.set_equilibrium(self.plasma.equilibrium) model = Interferometry(name="smmh1") model.set_los_transform(los_transform) @@ -246,7 +252,8 @@ def init_models(self): if "xrcs" in self.diagnostics: los_transform = self.ST40_data.binned_data["xrcs"]["te_kw"].transform model = Helike_spectroscopy(name="xrcs", window_masks=[slice(0.394, 0.396)], - window_vector=self.ST40_data.binned_data["xrcs"]["spectra"].wavelength.values * 0.1) + window_vector=self.ST40_data.binned_data["xrcs"][ + "spectra"].wavelength.values * 0.1) model.set_los_transform(los_transform) model.plasma = self.plasma self.models["xrcs"] = model @@ -264,48 +271,54 @@ def init_models(self): model.plasma = self.plasma self.models["cxff_pi"] = model - def phantom_data(self, noise=False, noise_factor=0.1): self.flat_data = {} if "smmh1" in self.diagnostics: - self.flat_data["smmh1.ne"] = self.models["smmh1"]().pop("ne").expand_dims(dim={"t": [self.plasma.time_to_calculate]}) + self.flat_data["smmh1.ne"] = self.models["smmh1"]().pop("ne").expand_dims( + dim={"t": [self.plasma.time_to_calculate]}) if "xrcs" in self.diagnostics: - self.flat_data["xrcs.spectra"] = self.models["xrcs"]().pop("spectra").expand_dims(dim={"t": [self.plasma.time_to_calculate]}) + self.flat_data["xrcs.spectra"] = self.models["xrcs"]().pop("spectra").expand_dims( + dim={"t": [self.plasma.time_to_calculate]}) self.flat_data["xrcs.background"] = None if "cxff_pi" in self.diagnostics: cxrs_data = self.models["cxff_pi"]().pop("ti").expand_dims(dim={"t": [self.plasma.time_to_calculate]}) self.flat_data["cxff_pi.ti"] = cxrs_data.where(cxrs_data.channel == 2) if "efit" in self.diagnostics: - self.flat_data["efit.wp"] = self.models["efit"]().pop("wp").expand_dims(dim={"t": [self.plasma.time_to_calculate]}) + self.flat_data["efit.wp"] = self.models["efit"]().pop("wp").expand_dims( + dim={"t": [self.plasma.time_to_calculate]}) if noise: - self.flat_data["smmh1.ne"] = self.flat_data["smmh1.ne"] + self.flat_data["smmh1.ne"].max().values * np.random.normal(0, noise_factor, None) - self.flat_data["xrcs.spectra"] = self.flat_data["xrcs.spectra"] + np.random.normal(0, np.sqrt(self.flat_data["xrcs.spectra"].values[0,]), self.flat_data["xrcs.spectra"].shape[1]) - self.flat_data["cxff_pi.ti"] = self.flat_data["cxff_pi.ti"] + self.flat_data["cxff_pi.ti"].max().values * np.random.normal(0, noise_factor, self.flat_data["cxff_pi.ti"].shape[1]) - self.flat_data["efit.wp"] = self.flat_data["efit.wp"] + self.flat_data["efit.wp"].max().values * np.random.normal(0, noise_factor, None) - + self.flat_data["smmh1.ne"] = self.flat_data["smmh1.ne"] + self.flat_data[ + "smmh1.ne"].max().values * np.random.normal(0, noise_factor, None) + self.flat_data["xrcs.spectra"] = self.flat_data["xrcs.spectra"] + np.random.normal(0, np.sqrt( + self.flat_data["xrcs.spectra"].values[0,]), self.flat_data["xrcs.spectra"].shape[1]) + self.flat_data["cxff_pi.ti"] = self.flat_data["cxff_pi.ti"] + self.flat_data[ + "cxff_pi.ti"].max().values * np.random.normal(0, noise_factor, self.flat_data["cxff_pi.ti"].shape[1]) + self.flat_data["efit.wp"] = self.flat_data["efit.wp"] + self.flat_data[ + "efit.wp"].max().values * np.random.normal(0, noise_factor, None) def exp_data(self): self.flat_data = flatdict.FlatDict(self.ST40_data.binned_data, ".") if "xrcs" in self.diagnostics: self.flat_data["xrcs.spectra"]["wavelength"] = self.flat_data["xrcs.spectra"].wavelength * 0.1 background = self.flat_data["xrcs.spectra"].where( - (self.flat_data["xrcs.spectra"].wavelength < 0.392) & - (self.flat_data["xrcs.spectra"].wavelength > 0.388), - drop=True) + (self.flat_data["xrcs.spectra"].wavelength < 0.392) & + (self.flat_data["xrcs.spectra"].wavelength > 0.388), + drop=True) self.flat_data["xrcs.background"] = background.mean(dim="wavelength") - self.flat_data["xrcs.spectra"]["error"] = np.sqrt(self.flat_data["xrcs.spectra"] + background.std(dim="wavelength") ** 2) + self.flat_data["xrcs.spectra"]["error"] = np.sqrt( + self.flat_data["xrcs.spectra"] + background.std(dim="wavelength") ** 2) if "cxff_pi" in self.diagnostics: - self.flat_data["cxff_pi"]["ti"] = self.flat_data["cxff_pi"]["ti"].where(self.flat_data["cxff_pi"]["ti"].channel==2, drop=True) - + self.flat_data["cxff_pi"]["ti"] = self.flat_data["cxff_pi"]["ti"].where( + self.flat_data["cxff_pi"]["ti"].channel == 2, drop=True) def __call__(self, *args, **kwargs): autocorr = sample_with_autocorr( self.sampler, self.start_points, iterations=self.iterations, auto_sample=5 ) - blobs = self.sampler.get_blobs(discard=self.burn_in, flat=True) + blobs = self.sampler.get_blobs(discard=int(self.iterations * self.burn_in_fraction), flat=True) blob_names = self.sampler.get_blobs().flatten()[0].keys() blob_dict = { blob_name: xr.concat( @@ -330,13 +343,13 @@ def __call__(self, *args, **kwargs): print(self.acceptance_fraction) plot_bayes_result(**result, figheader=self.result_path) -if __name__ == "__main__": - run = BayesWorkflow(pulse=10009, result_path="./results/10009_test/", iterations=1000, nwalkers=200, - burn_in=50, diagnostics=[ - "xrcs", - "efit", - "smmh1", - "cxff_pi" - ], phantom=False) +if __name__ == "__main__": + run = BayesWorkflow(pulse=10009, result_path="./results/10009_test/", iterations=100, nwalkers=50, + burn_in_fraction=0.10, diagnostics=[ + "xrcs", + "efit", + "smmh1", + "cxff_pi" + ], phantom=False) run() From 4b88ffe76b89969e601ff099fc25a4d9c94f0e24 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 23 May 2023 13:36:23 +0100 Subject: [PATCH 037/160] added violinplots --- indica/workflows/bayes_workflow.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/indica/workflows/bayes_workflow.py b/indica/workflows/bayes_workflow.py index b12e1986..08ed73e4 100644 --- a/indica/workflows/bayes_workflow.py +++ b/indica/workflows/bayes_workflow.py @@ -162,6 +162,23 @@ def _plot_1d( plt.savefig(figheader + filename) plt.close() +def violinplots(data, diag_data, filename): + nfig = len(data) + fig, axs = plt.subplots(1, nfig, figsize=(16, 6)) + for idx, key in enumerate(data.keys()): + axs[idx].violinplot(data[key]) + axs[idx].set_xlabel(key) + axs[idx].set_ylabel(data[key].datatype[0]) + axs[idx].axhline( + y=diag_data[key].sel(t=data[key].t).values, + color="black", + linestyle="-", + label=f"{key} data", + ) + + plt.setp([a.get_xticklabels() for a in axs], visible=False) + plt.savefig(filename) + plt.close() def plot_bayes_result( figheader="./results/test/", @@ -188,6 +205,12 @@ def plot_bayes_result( plt.savefig(figheader + "average_tau.png") plt.close() + violinplot_keys = ["efit.wp", "smmh1.ne", "cxff_pi.ti_0d"] + if "cxff_pi.ti" in blobs.keys(): + blobs["cxff_pi.ti_0d"] = blobs["cxff_pi.ti"].sel(channel = diag_data["cxff_pi.ti"].channel) + diag_data["cxff_pi.ti_0d"] = diag_data["cxff_pi.ti"].sel(channel=diag_data["cxff_pi.ti"].channel) + violinplots({key: blobs[key] for key in violinplot_keys if key in blobs.keys()}, diag_data, figheader+"boxplots.png") + key = "efit.wp" if key in blobs.keys(): _plot_0d( From 86e3784f9c9dfbc3f12cacf02d3b1c3912f0f1d2 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 23 May 2023 13:36:49 +0100 Subject: [PATCH 038/160] removed comment --- indica/bayesmodels.py | 1 - 1 file changed, 1 deletion(-) diff --git a/indica/bayesmodels.py b/indica/bayesmodels.py index 7d56c63e..4ef46d63 100644 --- a/indica/bayesmodels.py +++ b/indica/bayesmodels.py @@ -213,7 +213,6 @@ def ln_posterior(self, parameters: dict, **kwargs): t=self.plasma.time_to_calculate ), - # TODO: add Nh } blob = deepcopy({**self.bckc, **kin_profs}) return ln_posterior, blob From eb13c00d325de6f4c50d9edf188860e1994c0564 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 23 May 2023 13:37:53 +0100 Subject: [PATCH 039/160] included dictionary logic for fitting impurities scaled by concentrations and electron density --- indica/models/plasma.py | 64 ++++++++++++++++++----------------------- 1 file changed, 28 insertions(+), 36 deletions(-) diff --git a/indica/models/plasma.py b/indica/models/plasma.py index 83373f7b..18d0dc08 100644 --- a/indica/models/plasma.py +++ b/indica/models/plasma.py @@ -457,43 +457,32 @@ def initialize_variables(self, tstart: float, tend: float, dt: float): ], ) - def assign_profiles( - self, profile: str = "electron_density", t: float = None, element: str = "ar" - ): - if profile == "electron_density": - self.electron_density.loc[dict(t=t)] = self.Ne_prof() - self.calc_impurity_density( - tuple([imp for imp in self.impurities if imp != element])) # only update non-argon - elif profile == "electron_temperature": - self.electron_temperature.loc[dict(t=t)] = self.Te_prof() - elif profile == "ion_temperature": - self.ion_temperature.loc[dict(t=t)] = self.Ti_prof() - elif profile == "toroidal_rotation": - self.toroidal_rotation.loc[dict(t=t)] = self.Vrot_prof() - elif profile == "impurity_density": - self.impurity_density.loc[dict(t=t, element=element)] = self.Nimp_prof() - elif profile == "neutral_density": - self.neutral_density.loc[dict(t=t, element=element)] = self.Nh_prof() - else: - raise ValueError( - f"{profile} currently not found in possible Plasma properties" - ) + def assign_profiles(self, impurity_to_profile=None): + t = self.time_to_calculate + self.electron_density.loc[dict(t=t)] = self.Ne_prof() + self.calc_impurity_density(self.impurities) + self.electron_temperature.loc[dict(t=t)] = self.Te_prof() + self.ion_temperature.loc[dict(t=t)] = self.Ti_prof() + self.toroidal_rotation.loc[dict(t=t)] = self.Vrot_prof() + self.neutral_density.loc[dict(t=t, )] = self.Nh_prof() + if impurity_to_profile is not None: # overwrite impurity profile with non-Ne dependent profile + self.impurity_density.loc[dict(t=t, element=impurity_to_profile)] = self.Nimp_prof() def update_profiles( self, parameters: dict, - profile_prefixs: list = [ - "Te_prof", - "Ti_prof", - "Ne_prof", - "Nimp_prof", - "Vrot_prof", - ], ): """ Update plasma profiles with profile parameters i.e. {"Ne_prof.y0":1e19} -> Ne_prof.y0 """ + profile_prefixs: list = [ + "Te_prof", + "Ti_prof", + "Ne_prof", + "Nimp_prof", + "Vrot_prof", + ] for param, value in parameters.items(): _prefix = [pref for pref in profile_prefixs if pref in param] if _prefix: @@ -505,14 +494,17 @@ def update_profiles( else: raise ValueError(f"parameter: {key} not found in {prefix}") - for key in [ - "electron_density", - "electron_temperature", - "ion_temperature", - "toroidal_rotation", - "impurity_density", - ]: - self.assign_profiles(key, t=self.time_to_calculate) + imp_conc = [key.split("_")[0] for key in parameters.keys() if "conc" in key] + if any(imp_conc): + for idx, key in enumerate(self.impurities): + if any([imp for imp in imp_conc if key in imp]): + self.impurity_concentration[idx] = parameters[key+"_conc"] + + if any([key for key in parameters.keys() if "Nimp_prof" in key]): + impurity_to_profile = "ar" # TODO: generalise this for n independent impurity profiles e.g. Nimp1/Nimp2/... + else: + impurity_to_profile = None + self.assign_profiles(impurity_to_profile=impurity_to_profile) @property def time_to_calculate(self): From 898731150ac91f3ecfb200895b4c4a26570333d9 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 23 May 2023 15:47:43 +0100 Subject: [PATCH 040/160] changed name cxff_pi.ti_0d to cxff_pi.ti0 --- indica/workflows/bayes_workflow.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/indica/workflows/bayes_workflow.py b/indica/workflows/bayes_workflow.py index 08ed73e4..ce82b8d1 100644 --- a/indica/workflows/bayes_workflow.py +++ b/indica/workflows/bayes_workflow.py @@ -193,6 +193,7 @@ def plot_bayes_result( ): Path(figheader).mkdir(parents=True, exist_ok=True) + plt.figure() plt.plot( np.arange(0, autocorr.__len__())[np.isfinite(autocorr)], @@ -205,10 +206,10 @@ def plot_bayes_result( plt.savefig(figheader + "average_tau.png") plt.close() - violinplot_keys = ["efit.wp", "smmh1.ne", "cxff_pi.ti_0d"] + violinplot_keys = ["efit.wp", "smmh1.ne", "cxff_pi.ti0"] if "cxff_pi.ti" in blobs.keys(): - blobs["cxff_pi.ti_0d"] = blobs["cxff_pi.ti"].sel(channel = diag_data["cxff_pi.ti"].channel) - diag_data["cxff_pi.ti_0d"] = diag_data["cxff_pi.ti"].sel(channel=diag_data["cxff_pi.ti"].channel) + blobs["cxff_pi.ti0"] = blobs["cxff_pi.ti"].sel(channel = diag_data["cxff_pi.ti"].channel) + diag_data["cxff_pi.ti0"] = diag_data["cxff_pi.ti"].sel(channel=diag_data["cxff_pi.ti"].channel) violinplots({key: blobs[key] for key in violinplot_keys if key in blobs.keys()}, diag_data, figheader+"boxplots.png") key = "efit.wp" From a9b6b18000d06958c29832075c2db2990fd7198d Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Wed, 24 May 2023 14:18:49 +0100 Subject: [PATCH 041/160] new option for initialising at maximum-likelihood estimate --- indica/workflows/bayes_dev_workflow.py | 58 +++++++++++++++++++------- 1 file changed, 43 insertions(+), 15 deletions(-) diff --git a/indica/workflows/bayes_dev_workflow.py b/indica/workflows/bayes_dev_workflow.py index b7a094a0..cb25407a 100644 --- a/indica/workflows/bayes_dev_workflow.py +++ b/indica/workflows/bayes_dev_workflow.py @@ -53,7 +53,7 @@ "ar_conc": loguniform(0.0001, 0.01), - "Nimp_prof.y0": get_uniform(1e15, 1e17), + "Nimp_prof.y0": get_uniform(1e15, 1e18), "Nimp_prof.y1": get_uniform(1e15, 2e16), "Ne_prof.y0/Nimp_prof.y0": lambda x1, x2: np.where((x1 > x2 * 100) & (x1 < x2 * 1e5), 1, 0), "Nimp_prof.y0/Nimp_prof.y1": lambda x1, x2: np.where((x1 > x2), 1, 0), @@ -75,28 +75,28 @@ } OPTIMISED_PARAMS = [ - "Ne_prof.y0", # "Ne_prof.y1", + "Ne_prof.y0", "Ne_prof.peaking", "Ne_prof.wcenter", "Ne_prof.wped", # "ar_conc", + # "Nimp_prof.y1", "Nimp_prof.y0", "Nimp_prof.wcenter", "Nimp_prof.wped", - # "Nimp_prof.y1", "Nimp_prof.peaking", "Te_prof.y0", - # "Te_prof.wped", - # "Te_prof.wcenter", - # "Te_prof.peaking", + "Te_prof.wped", + "Te_prof.wcenter", + "Te_prof.peaking", "Ti_prof.y0", - # "Ti_prof.wped", - # "Ti_prof.wcenter", - # "Ti_prof.peaking", + "Ti_prof.wped", + "Ti_prof.wcenter", + "Ti_prof.peaking", ] OPTIMISED_QUANTITY = [ @@ -118,9 +118,10 @@ def __init__(self, tsample=4, iterations=100, burn_in_fraction=0, + center_mass_sampling = True, diagnostics=None, + profiles=None, phantom=True, - profiles=None ): self.pulse = pulse @@ -132,6 +133,8 @@ def __init__(self, self.nwalkers = nwalkers self.iterations = iterations self.burn_in_fraction = burn_in_fraction + self.center_mass_sampling = center_mass_sampling + self.diagnostics = diagnostics self.phantom = phantom self.profiles = profiles @@ -169,9 +172,7 @@ def __init__(self, ) ndim = len(OPTIMISED_PARAMS) - self.start_points = self.bayes_run.sample_from_priors(OPTIMISED_PARAMS, size=self.nwalkers) self.move = [(emcee.moves.StretchMove(), 0.9), (emcee.moves.DEMove(), 0.1)] - self.sampler = emcee.EnsembleSampler( nwalkers, ndim, @@ -182,6 +183,33 @@ def __init__(self, "background": self.flat_data["xrcs.background"]}, ) + if not self.center_mass_sampling: + self.start_points = self.bayes_run.sample_from_priors(OPTIMISED_PARAMS, size=self.nwalkers) + else: + # gaussian mass around most probable starting points + nsamples = 100 + start_points = self.bayes_run.sample_from_priors(OPTIMISED_PARAMS, size=nsamples) + ln_prob, _ = self.sampler.compute_log_prob(start_points) + num_best_points = int(nsamples * 0.05) + index_best_start = np.argsort(ln_prob)[-num_best_points:] + best_start_points = start_points[index_best_start,:] + best_points_std = np.std(best_start_points, axis=0) + + # Passing samples through ln_prior and redrawing if they fail + samples = np.empty((OPTIMISED_PARAMS.__len__(), 0)) + while samples.size < OPTIMISED_PARAMS.__len__() * self.nwalkers: + sample = np.random.normal(np.mean(best_start_points, axis=0), best_points_std * 2, + size=(self.nwalkers*5, len(OPTIMISED_PARAMS))) + start = {name: sample[:, idx] for idx, name in enumerate(OPTIMISED_PARAMS)} + ln_prior = self.bayes_run._ln_prior(start) + # Convert from dictionary of arrays -> array, + # then filtering out where ln_prior is -infinity + accepted_samples = np.array(list(start.values()))[ + :, ln_prior != -np.inf + ] + samples = np.append(samples, accepted_samples, axis=1) + self.start_points = samples[:, 0:self.nwalkers].transpose() + def init_fast_particles(self): st40_code = ReadST40(13110009, self.tstart, self.tend, dt=self.dt, tree="astra") st40_code.get_raw_data("", "astra", "RUN573") @@ -345,11 +373,11 @@ def __call__(self, *args, **kwargs): if __name__ == "__main__": - run = BayesWorkflow(pulse=10009, result_path="./results/10009_test/", iterations=100, nwalkers=50, - burn_in_fraction=0.10, diagnostics=[ + run = BayesWorkflow(pulse=10009, result_path="./results/10009_test/", iterations=5000, nwalkers=100, + burn_in_fraction=0.20, diagnostics=[ "xrcs", "efit", "smmh1", "cxff_pi" - ], phantom=False) + ], phantom=False, center_mass_sampling=True) run() From 1fa84279ac190e8054092208837029466e0fd7e8 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Wed, 24 May 2023 14:19:40 +0100 Subject: [PATCH 042/160] added xlim for violin plot --- indica/workflows/bayes_workflow.py | 1 + 1 file changed, 1 insertion(+) diff --git a/indica/workflows/bayes_workflow.py b/indica/workflows/bayes_workflow.py index ce82b8d1..25de08be 100644 --- a/indica/workflows/bayes_workflow.py +++ b/indica/workflows/bayes_workflow.py @@ -168,6 +168,7 @@ def violinplots(data, diag_data, filename): for idx, key in enumerate(data.keys()): axs[idx].violinplot(data[key]) axs[idx].set_xlabel(key) + axs[idx].set_ylim(bottom=0) axs[idx].set_ylabel(data[key].datatype[0]) axs[idx].axhline( y=diag_data[key].sel(t=data[key].t).values, From db2d0e1f8d77fb4bb239312d7ec31a1ade254140 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Wed, 31 May 2023 13:48:51 +0100 Subject: [PATCH 043/160] tweaking calibration --- indica/models/helike_spectroscopy.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/indica/models/helike_spectroscopy.py b/indica/models/helike_spectroscopy.py index eae06066..a0d3aade 100644 --- a/indica/models/helike_spectroscopy.py +++ b/indica/models/helike_spectroscopy.py @@ -40,7 +40,7 @@ def __init__( name: str, instrument_method="get_helike_spectroscopy", etendue: float = 1.0, - calibration: float = 4.0e-20, + calibration: float = 8.0e-20, marchuk: bool = True, full_run: bool = False, element: str = "ar", From ce73ee5b823c3a8c3c33a44d9122e1fe42fdfb1e Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Wed, 31 May 2023 13:50:43 +0100 Subject: [PATCH 044/160] fixed neutral density units --- indica/models/plasma.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/indica/models/plasma.py b/indica/models/plasma.py index 18d0dc08..33014646 100644 --- a/indica/models/plasma.py +++ b/indica/models/plasma.py @@ -307,7 +307,7 @@ def initialize_variables(self, tstart: float, tend: float, dt: float): self.data2d, ("density", "electron"), "$m^{-3}$" ) self.neutral_density: DataArray = assign_data( - self.data2d, ("thermal_neutral", "density"), "eV" + self.data2d, ("density", "thermal_neutral"), "$m^{-3}$" ) self.tau: DataArray = assign_data(self.data2d, ("time", "residence"), "s") @@ -462,7 +462,10 @@ def assign_profiles(self, impurity_to_profile=None): self.electron_density.loc[dict(t=t)] = self.Ne_prof() self.calc_impurity_density(self.impurities) self.electron_temperature.loc[dict(t=t)] = self.Te_prof() - self.ion_temperature.loc[dict(t=t)] = self.Ti_prof() + if self.Ti_ref: + self.ion_temperature.loc[dict(t=t)] = self.Ti_prof(y0_ref=self.Te_prof.y0) + else: + self.ion_temperature.loc[dict(t=t)] = self.Ti_prof() self.toroidal_rotation.loc[dict(t=t)] = self.Vrot_prof() self.neutral_density.loc[dict(t=t, )] = self.Nh_prof() if impurity_to_profile is not None: # overwrite impurity profile with non-Ne dependent profile From eda445ae0bd9cfdb576df5aa828f1a0942a06723 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Wed, 31 May 2023 13:51:20 +0100 Subject: [PATCH 045/160] added set plot functions --- indica/utilities.py | 95 ++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 94 insertions(+), 1 deletion(-) diff --git a/indica/utilities.py b/indica/utilities.py index 585334eb..b8148516 100644 --- a/indica/utilities.py +++ b/indica/utilities.py @@ -1,7 +1,7 @@ """Various miscellanious helper functions.""" - from copy import deepcopy import inspect +import os import string from typing import Any from typing import Callable @@ -12,6 +12,9 @@ from typing import Tuple from typing import Union +from matplotlib import cm +from matplotlib import rcParams +import matplotlib.pylab as plt import numpy as np from scipy.interpolate import CubicSpline from xarray import apply_ufunc @@ -398,3 +401,93 @@ def check_time_present(t_desired: LabeledArray, t_array: LabeledArray): ) if not equil_ok: raise ValueError(f"Desired time {t_desired} not available in array {t_array}") + + +def save_figure( + path_name: str = "", + fig_name: str = "", + orientation: str = "landscape", + dpi: int = 300, + quality: int = 95, + ext: str = "png", + save_fig: bool = True, + create_path: bool = True, +): + if save_fig: + _fig_name = deepcopy(fig_name) + _path_name = deepcopy(path_name) + if _path_name[-1] != "/": + _path_name = f"{_path_name}/" + _file = f"{_path_name}{_fig_name}.{ext}" + + kwargs = {"orientation": orientation, "dpi": dpi} + if ext != "svg": + kwargs["pil_kwargs"] = {"quality": quality} + + if create_path and not os.path.exists(_path_name): + os.mkdir(_path_name) + print(f"Creating directory {_path_name}") + + plt.savefig( + _file, + **kwargs, + ) + print(f"Saving picture to {_file}") + + +def set_plot_colors( + color_map: str = "gnuplot2", +): + + cmap = getattr(cm, color_map) + colors = { + "electron": cmap(0.1), + "ion": cmap(0.75), + "fast": cmap(0.4), + "impurity": cmap(0.0), + "raw": "black", + "thermal": "red", + "total": "black", + "binned": cmap(0.2), + "bckc": cmap(0.75), + } + + return cmap, colors + + +def set_plot_rcparams(option: str = "profiles"): + plot_params: dict = { + "profiles": { + "font.size": 13, + "legend.fontsize": 11, + "lines.markersize": 10, + "lines.linewidth": 2, + }, + "multi": { + "font.size": 13, + "legend.fontsize": 13, + "lines.markersize": 10, + "lines.linewidth": 2, + "figure.figsize": [6.4, 3.8], + "figure.subplot.bottom": 0.15, + }, + "time_evolution": { + "font.size": 11, + "legend.fontsize": 8, + "lines.markersize": 5, + "lines.linewidth": 2, + "font.weight": 600, + }, + } + + if option not in plot_params.keys(): + return + + for key, value in plot_params[option].items(): + rcParams.update({key: value}) + + +def set_axis_sci(axis: str = "y"): + ylim = np.abs(np.array(plt.ylim())) + if ylim.max() > 1.0e3 or ylim.min() < 1.0e-2: + plt.ticklabel_format(style="sci", axis=axis, scilimits=(0, 0)) From 2555605f6db19eef8bd2cbe7e0c7a57fc07ebdbd Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Wed, 31 May 2023 13:52:39 +0100 Subject: [PATCH 046/160] added center mass sampling and Ti/Te y0_ref options --- indica/workflows/bayes_dev_workflow.py | 41 +++++++++++++++----------- 1 file changed, 24 insertions(+), 17 deletions(-) diff --git a/indica/workflows/bayes_dev_workflow.py b/indica/workflows/bayes_dev_workflow.py index cb25407a..ac0110d0 100644 --- a/indica/workflows/bayes_dev_workflow.py +++ b/indica/workflows/bayes_dev_workflow.py @@ -33,13 +33,15 @@ "Nimp_prof.peaking": 2, "Te_prof.y0": 3000, + "Te_prof.y1": 50, "Te_prof.wcenter": 0.4, - "Te_prof.wped": 4, + "Te_prof.wped": 3, "Te_prof.peaking": 2, "Ti_prof.y0": 6000, + "Ti_prof.y1": 50, "Ti_prof.wcenter": 0.4, - "Ti_prof.wped": 4, + "Ti_prof.wped": 3, "Ti_prof.peaking": 2, } @@ -53,13 +55,13 @@ "ar_conc": loguniform(0.0001, 0.01), - "Nimp_prof.y0": get_uniform(1e15, 1e18), + "Nimp_prof.y0": get_uniform(1e16, 1e18), "Nimp_prof.y1": get_uniform(1e15, 2e16), "Ne_prof.y0/Nimp_prof.y0": lambda x1, x2: np.where((x1 > x2 * 100) & (x1 < x2 * 1e5), 1, 0), "Nimp_prof.y0/Nimp_prof.y1": lambda x1, x2: np.where((x1 > x2), 1, 0), "Nimp_prof.wped": get_uniform(1, 6), "Nimp_prof.wcenter": get_uniform(0.1, 0.8), - "Nimp_prof.peaking": get_uniform(1, 6), + "Nimp_prof.peaking": get_uniform(1, 20), "Nimp_prof.peaking/Ne_prof.peaking": lambda x1, x2: np.where((x1 > x2), 1, 0), # impurity always more peaked "Te_prof.y0": get_uniform(1000, 5000), @@ -68,10 +70,10 @@ "Te_prof.peaking": get_uniform(1, 6), "Ti_prof.y0/Te_prof.y0": lambda x1, x2: np.where(x1 > x2, 1, 0), # hot ion mode - "Ti_prof.y0": get_uniform(2000, 10000), + "Ti_prof.y0": get_uniform(3000, 10000), "Ti_prof.wped": get_uniform(1, 6), "Ti_prof.wcenter": get_uniform(0.1, 0.6), - "Ti_prof.peaking": get_uniform(1, 6), + "Ti_prof.peaking": get_uniform(1, 20), } OPTIMISED_PARAMS = [ @@ -115,10 +117,11 @@ def __init__(self, tstart=0.02, tend=0.10, dt=0.01, - tsample=4, + tsample=0.06, iterations=100, burn_in_fraction=0, - center_mass_sampling = True, + center_mass_sampling=True, + Ti_ref=True, diagnostics=None, profiles=None, phantom=True, @@ -134,6 +137,7 @@ def __init__(self, self.iterations = iterations self.burn_in_fraction = burn_in_fraction self.center_mass_sampling = center_mass_sampling + self.Ti_ref = Ti_ref self.diagnostics = diagnostics self.phantom = phantom @@ -149,7 +153,9 @@ def __init__(self, full_run=False, n_rad=20, ) - self.plasma.time_to_calculate = self.plasma.t[tsample] + self.plasma.Ti_ref = self.Ti_ref + self.tsample = tsample + self.plasma.time_to_calculate = self.plasma.t[np.abs(tsample-self.plasma.t).argmin()] self.plasma.update_profiles(DEFAULT_PHANTOM_PARAMS) self.plasma.build_atomic_data(calc_power_loss=False) @@ -192,14 +198,14 @@ def __init__(self, ln_prob, _ = self.sampler.compute_log_prob(start_points) num_best_points = int(nsamples * 0.05) index_best_start = np.argsort(ln_prob)[-num_best_points:] - best_start_points = start_points[index_best_start,:] + best_start_points = start_points[index_best_start, :] best_points_std = np.std(best_start_points, axis=0) # Passing samples through ln_prior and redrawing if they fail samples = np.empty((OPTIMISED_PARAMS.__len__(), 0)) while samples.size < OPTIMISED_PARAMS.__len__() * self.nwalkers: sample = np.random.normal(np.mean(best_start_points, axis=0), best_points_std * 2, - size=(self.nwalkers*5, len(OPTIMISED_PARAMS))) + size=(self.nwalkers * 5, len(OPTIMISED_PARAMS))) start = {name: sample[:, idx] for idx, name in enumerate(OPTIMISED_PARAMS)} ln_prior = self.bayes_run._ln_prior(start) # Convert from dictionary of arrays -> array, @@ -225,7 +231,7 @@ def init_fast_particles(self): self.plasma.pressure_fast_perpendicular.values = Pbper.values def read_st40(self, diagnostics=None): - self.ST40_data = ReadST40(self.pulse, tstart=self.tstart, tend=self.tend) + self.ST40_data = ReadST40(self.pulse, tstart=self.tstart, tend=self.tend, dt=self.dt) self.ST40_data(diagnostics) self.plasma.set_equilibrium(self.ST40_data.equilibrium) @@ -357,7 +363,7 @@ def __call__(self, *args, **kwargs): } samples = self.sampler.get_chain(flat=True) - prior_samples = self.bayes_run.sample_from_priors(OPTIMISED_PARAMS, size=int(1e5)) + prior_samples = self.bayes_run.sample_from_priors(OPTIMISED_PARAMS, size=int(1e4)) result = { "blobs": blob_dict, "diag_data": self.flat_data, @@ -369,15 +375,16 @@ def __call__(self, *args, **kwargs): } self.acceptance_fraction = self.sampler.acceptance_fraction.sum() print(self.acceptance_fraction) - plot_bayes_result(**result, figheader=self.result_path) + plot_bayes_result(result, figheader=self.result_path, filetype=".png") if __name__ == "__main__": - run = BayesWorkflow(pulse=10009, result_path="./results/10009_test/", iterations=5000, nwalkers=100, - burn_in_fraction=0.20, diagnostics=[ + run = BayesWorkflow(pulse=10009, result_path="./results/10009_60ms_long/", iterations=500, nwalkers=200, + burn_in_fraction=0.10, dt=0.005, tsample=0.060, + diagnostics=[ "xrcs", "efit", "smmh1", "cxff_pi" - ], phantom=False, center_mass_sampling=True) + ], phantom=False, center_mass_sampling=True, Ti_ref=True) run() From 20ae704ee7e8df3e43f6f99587731e9b2de7b56c Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Wed, 31 May 2023 13:53:07 +0100 Subject: [PATCH 047/160] plot formatting changed --- indica/workflows/bayes_workflow.py | 270 ++++++++++++++++++----------- 1 file changed, 168 insertions(+), 102 deletions(-) diff --git a/indica/workflows/bayes_workflow.py b/indica/workflows/bayes_workflow.py index 25de08be..e29add98 100644 --- a/indica/workflows/bayes_workflow.py +++ b/indica/workflows/bayes_workflow.py @@ -3,26 +3,46 @@ import corner import matplotlib.pyplot as plt import numpy as np +import pickle + +from indica.utilities import set_plot_rcparams, set_axis_sci def plot_profile( - profile, - blobkey: str, - figheader="./results/test/", - phantom_profile=None, - sharefig=False, - filename="", - linestyle="--", - color="blue", + profile, + blobkey: str, + figheader="./results/test/", + phantom_profile=None, + logscale=False, + sharefig=False, + filename="", + filetype=".png", + linestyle="--", + color="blue", ): - if not plt.get_fignums(): # If no figure is open - plt.figure(figsize=(8, 6)) + set_plot_rcparams("profiles") + # if not plt.get_fignums(): # If no figure is open + # plt.figure(figsize=(8, 6)) + if blobkey == "electron_temperature": + legkey = "Te" + elif blobkey == "ion_temperature": + legkey = "Ti" + elif blobkey == "ion_density": + legkey = "Ni" + elif blobkey == "electron_density": + legkey = "Ne" + elif blobkey == "impurity_density": + legkey = "Nimp" + elif blobkey == "fast_density": + legkey = "Nfast" + elif blobkey == "neutral_density": + legkey = "Nneut" plt.fill_between( profile.rho_poloidal, profile.quantile(0.16, dim="index"), profile.quantile(0.84, dim="index"), - label=f"{blobkey}, 68% Confidence", + label=f"{legkey}, 68% Confidence", zorder=3, color=color, alpha=0.8, @@ -31,7 +51,7 @@ def plot_profile( profile.rho_poloidal, profile.quantile(0.025, dim="index"), profile.quantile(0.975, dim="index"), - label=f"{blobkey}, 95% Confidence", + label=f"{legkey}, 95% Confidence", zorder=2, color="grey", alpha=0.6, @@ -40,7 +60,7 @@ def plot_profile( profile.rho_poloidal, profile.quantile(0.00, dim="index"), profile.quantile(1.00, dim="index"), - label=f"{blobkey}, Max-Min", + label=f"{legkey}, Max-Min", zorder=1, color="lightgrey", alpha=0.6, @@ -52,7 +72,7 @@ def plot_profile( else: phantom = phantom_profile[blobkey] phantom.plot( - label=f"{blobkey}, phantom profile", + label=f"{legkey}, phantom profile", linestyle=linestyle, color="black", zorder=4, @@ -62,52 +82,30 @@ def plot_profile( if sharefig: return + set_axis_sci() + if logscale: + plt.yscale("log") + + plt.xlabel("Rho poloidal") + plt.ylabel(f"{profile.datatype[0].capitalize()} [{profile.units}]") if filename: - plt.savefig(figheader + f"{filename}.png") + plt.savefig(figheader + f"{filename}{filetype}") else: - plt.savefig(figheader + f"{blobkey}.png") + plt.savefig(figheader + f"{blobkey}{filetype}") plt.close("all") -def _plot_0d( - blobs: dict, - blobkey: str, - diag_data: dict, - filename: str, - figheader="./results/test/", - xlabel="samples ()", - ylabel="a.u.", - figsize=(6.4, 4.8), - **kwargs, -): - if blobkey not in blobs.keys(): - raise ValueError(f"{blobkey} not in blobs") - plt.figure(figsize=figsize) - blob_data = blobs[blobkey] - plt.plot(blob_data, label=f"{blobkey} model") - plt.axhline( - y=diag_data[blobkey].sel(t=blob_data.t).values, - color="black", - linestyle="-", - label=f"{blobkey} data", - ) - plt.xlabel(xlabel) - plt.ylabel(ylabel) - plt.legend() - plt.savefig(figheader + filename) - plt.close() - - def _plot_1d( - blobs: dict, - blobkey: str, - diag_data: dict, - filename: str, - figheader="./results/test/", - ylabel="a.u.", - xlim = None, - figsize=(6.4, 4.8), - **kwargs, + blobs: dict, + blobkey: str, + diag_data: dict, + filename: str, + figheader="./results/test/", + ylabel="a.u.", + xlabel="[]", + xlim=None, + figsize=(6.4, 4.8), + **kwargs, ): if blobkey not in blobs.keys(): raise ValueError(f"{blobkey} not in blobs") @@ -155,45 +153,86 @@ def _plot_1d( label=f"{blobkey} data", zorder=4, ) + set_axis_sci() plt.ylabel(ylabel) - plt.xlabel(dims[0]) + plt.xlabel(xlabel) plt.xlim(xlim) plt.legend() plt.savefig(figheader + filename) plt.close() -def violinplots(data, diag_data, filename): + +def violinplot(data: dict, key: str, diag_data: dict, filename: str, ylabel="[a.u.]", figheader="./results/test/", + **kwargs): + set_plot_rcparams() + fig, axs = plt.subplots(1, 1,) + _data = data[key].values + _data = _data[((_data > np.quantile(_data, 0.16)) & (_data < np.quantile(_data, 0.84)))] + violin = axs.violinplot(_data, + showextrema=False, + # quantiles=[0.025, 0.975, 0.16, 0.84], + # showmedians=True, + ) + violin["bodies"][0].set_edgecolor("black") + axs.set_xlabel(key) + top = axs.get_ylim()[1] + axs.set_ylim(bottom=0, top=top * 1.1) + axs.set_ylabel(f"{ylabel}") + y = diag_data[key].sel(t=data[key].t).values + axs.errorbar(1, y=y, yerr=y * 0.10, fmt="D", ecolor="black", capsize=10, color="black") + set_axis_sci() + plt.setp([axs.get_xticklabels()], visible=False) + plt.savefig(figheader + filename) + plt.close() + + +def histograms(data, diag_data, filename): nfig = len(data) fig, axs = plt.subplots(1, nfig, figsize=(16, 6)) for idx, key in enumerate(data.keys()): - axs[idx].violinplot(data[key]) - axs[idx].set_xlabel(key) - axs[idx].set_ylim(bottom=0) - axs[idx].set_ylabel(data[key].datatype[0]) - axs[idx].axhline( - y=diag_data[key].sel(t=data[key].t).values, + n, bins, patches = axs[idx].hist(data[key], 50, density=True) + q1 = (np.percentile(data[key], 16), np.percentile(data[key], 84)) + q2 = (np.percentile(data[key], 2.5), np.percentile(data[key], 97.5)) + idx_high = np.argwhere((bins > q1[0]) & (bins < q1[1])).flatten() + idx_low = np.argwhere((bins > q2[0]) & (bins < q2[1])).flatten() + for patch in patches: + patch.set_facecolor("lightgrey") + for i in idx_low: + patches[i].set_facecolor("grey") + for i in idx_high: + patches[i].set_facecolor("red") + + axs[idx].set_xlabel(f"{key} ({data[key].datatype[0]})") + + axs[idx].axvline( + x=diag_data[key].sel(t=data[key].t).values, color="black", - linestyle="-", + linestyle="-.", label=f"{key} data", ) + axs[0].set_ylabel("pdf ()") - plt.setp([a.get_xticklabels() for a in axs], visible=False) plt.savefig(filename) plt.close() + def plot_bayes_result( - figheader="./results/test/", - blobs=None, - diag_data=None, - samples=None, - prior_samples=None, - param_names=None, - phantom_profiles=None, - autocorr=None, - **kwargs, + results, + figheader="./results/test/", + filetype=".png", + **kwargs, ): - Path(figheader).mkdir(parents=True, exist_ok=True) + diag_data = results["diag_data"] + blobs = results["blobs"] + samples = results["samples"] + prior_samples = results["prior_samples"] + param_names = results["param_names"] + phantom_profiles = results["phantom_profiles"] + autocorr = results["autocorr"] + Path(figheader).mkdir(parents=True, exist_ok=True) + with open(figheader + "results.pkl", "wb") as handle: + pickle.dump(results, handle) plt.figure() plt.plot( @@ -204,54 +243,63 @@ def plot_bayes_result( plt.legend() plt.xlabel("iterations") plt.ylabel("auto-correlation time (iterations)") - plt.savefig(figheader + "average_tau.png") + plt.savefig(figheader + "average_tau" + filetype) plt.close() - violinplot_keys = ["efit.wp", "smmh1.ne", "cxff_pi.ti0"] if "cxff_pi.ti" in blobs.keys(): - blobs["cxff_pi.ti0"] = blobs["cxff_pi.ti"].sel(channel = diag_data["cxff_pi.ti"].channel) + blobs["cxff_pi.ti0"] = blobs["cxff_pi.ti"].sel(channel=diag_data["cxff_pi.ti"].channel) diag_data["cxff_pi.ti0"] = diag_data["cxff_pi.ti"].sel(channel=diag_data["cxff_pi.ti"].channel) - violinplots({key: blobs[key] for key in violinplot_keys if key in blobs.keys()}, diag_data, figheader+"boxplots.png") + + key = "cxff_pi.ti0" + if key in blobs.keys(): + violinplot( + blobs, + key, + diag_data, + f"{key.replace('.', '_')}" + filetype, + figheader=figheader, + ylabel="Temperature [eV]", + ) key = "efit.wp" if key in blobs.keys(): - _plot_0d( + violinplot( blobs, key, diag_data, - f"{key.replace('.', '_')}.png", + f"{key.replace('.', '_')}" + filetype, figheader=figheader, - ylabel="Wp (J)", + ylabel="Energy [J]", ) key = "smmh1.ne" if key in blobs.keys(): - _plot_0d( + violinplot( blobs, key, diag_data, - f"{key.replace('.', '_')}.png", + f"{key.replace('.', '_')}" + filetype, figheader=figheader, - ylabel="ne_int (m^-2)", + ylabel=r"Line Integrated Density [$m^{-2}$]", ) key = "xrcs.te_kw" if key in blobs.keys(): - _plot_0d( + violinplot( blobs, key, diag_data, - f"{key.replace('.', '_')}.png", + f"{key.replace('.', '_')}" + filetype, figheader=figheader, - ylabel="temperature (eV)", + ylabel="Temperature [eV]", ) key = "xrcs.ti_w" if key in blobs.keys(): - _plot_0d( + violinplot( blobs, key, diag_data, - f"{key.replace('.', '_')}.png", + f"{key.replace('.', '_')}" + filetype, figheader=figheader, - ylabel="temperature (eV)", + ylabel="Temperature [eV]", ) key = "xrcs.spectra" if key in blobs.keys(): @@ -259,11 +307,12 @@ def plot_bayes_result( blobs, key, diag_data, - f"{key.replace('.', '_')}.png", + f"{key.replace('.', '_')}" + filetype, figheader=figheader, - ylabel="intensity (A.U.)", - xlim = (0.394, 0.401), - figsize=(12, 10), + ylabel="Intensity [a.u.]", + xlabel="Wavelength [nm]", + xlim=(0.394, 0.401), + figsize=(6, 4), ) key = "cxff_pi.ti" if key in blobs.keys(): @@ -271,9 +320,11 @@ def plot_bayes_result( blobs, key, diag_data, - f"{key.replace('.', '_')}.png", + f"{key.replace('.', '_')}" + filetype, figheader=figheader, - ylabel="temperature (eV)", + ylabel="Temperature [eV]", + xlabel="Channel", + ) key = "electron_temperature" @@ -281,6 +332,7 @@ def plot_bayes_result( blobs[key], key, figheader=figheader, + filetype=filetype, phantom_profile=phantom_profiles, sharefig=True, color="blue", @@ -292,13 +344,14 @@ def plot_bayes_result( key, figheader=figheader, filename="temperature", + filetype=filetype, phantom_profile=phantom_profiles, color="red", linestyle="dotted", ) key = "electron_density" plot_profile( - blobs[key], key, figheader=figheader, phantom_profile=phantom_profiles + blobs[key], key, figheader=figheader, filetype=filetype, phantom_profile=phantom_profiles ) key = "impurity_density" for elem in blobs[key].element.values: @@ -307,6 +360,7 @@ def plot_bayes_result( key, figheader=figheader, filename=f"{elem} density", + filetype=filetype, phantom_profile=phantom_profiles, color="red", ) @@ -316,6 +370,7 @@ def plot_bayes_result( key, figheader=figheader, filename=f"h density", + filetype=filetype, phantom_profile=phantom_profiles, color="red", ) @@ -324,6 +379,7 @@ def plot_bayes_result( blobs[key], key, figheader=figheader, + filetype=filetype, phantom_profile=phantom_profiles, color="red", ) @@ -331,15 +387,18 @@ def plot_bayes_result( plot_profile( blobs[key], key, + filename="", figheader=figheader, + filetype=filetype, phantom_profile=phantom_profiles, + logscale=True ) corner.corner(samples, labels=param_names) - plt.savefig(figheader + "posterior.png") + plt.savefig(figheader + "posterior" + filetype) corner.corner(prior_samples, labels=param_names) - plt.savefig(figheader + "prior.png") + plt.savefig(figheader + "prior" + filetype) plt.close("all") @@ -347,9 +406,9 @@ def sample_with_autocorr(sampler, start_points, iterations=10, auto_sample=5): autocorr = np.ones((iterations,)) * np.nan old_tau = np.inf for sample in sampler.sample( - start_points, - iterations=iterations, - progress=True, + start_points, + iterations=iterations, + progress=True, ): if sampler.iteration % auto_sample: continue @@ -362,3 +421,10 @@ def sample_with_autocorr(sampler, start_points, iterations=10, auto_sample=5): old_tau = new_tau autocorr = autocorr[: sampler.iteration] return autocorr + + +if __name__ == "__main__": + filehead = "./results/10009_60ms_short/" + with open(filehead + "results.pkl", "rb") as handle: + results = pickle.load(handle) + plot_bayes_result(results, filehead, filetype=".png") From e0b9c03518477813d2a060f42f288c18af984346 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Wed, 31 May 2023 13:57:44 +0100 Subject: [PATCH 048/160] black formatting --- indica/models/plasma.py | 164 +++++++++--------- indica/workflows/bayes_dev_workflow.py | 227 ++++++++++++++++--------- indica/workflows/bayes_workflow.py | 106 +++++++----- 3 files changed, 294 insertions(+), 203 deletions(-) diff --git a/indica/models/plasma.py b/indica/models/plasma.py index 33014646..1f57439c 100644 --- a/indica/models/plasma.py +++ b/indica/models/plasma.py @@ -103,18 +103,18 @@ class Plasma: def __init__( - self, - tstart: float = 0.01, - tend: float = 0.14, - dt: float = 0.01, - machine_dimensions=((0.15, 0.95), (-0.7, 0.7)), - impurities: tuple = ("c", "ar"), - main_ion: str = "h", - impurity_concentration: tuple = (0.02, 0.001), - pulse: int = None, - full_run: bool = False, - n_rad: int = 41, - verbose: bool = False, + self, + tstart: float = 0.01, + tend: float = 0.14, + dt: float = 0.01, + machine_dimensions=((0.15, 0.95), (-0.7, 0.7)), + impurities: tuple = ("c", "ar"), + main_ion: str = "h", + impurity_concentration: tuple = (0.02, 0.001), + pulse: int = None, + full_run: bool = False, + n_rad: int = 41, + verbose: bool = False, ): """ Class for plasma objects. @@ -467,13 +467,21 @@ def assign_profiles(self, impurity_to_profile=None): else: self.ion_temperature.loc[dict(t=t)] = self.Ti_prof() self.toroidal_rotation.loc[dict(t=t)] = self.Vrot_prof() - self.neutral_density.loc[dict(t=t, )] = self.Nh_prof() - if impurity_to_profile is not None: # overwrite impurity profile with non-Ne dependent profile - self.impurity_density.loc[dict(t=t, element=impurity_to_profile)] = self.Nimp_prof() + self.neutral_density.loc[ + dict( + t=t, + ) + ] = self.Nh_prof() + if ( + impurity_to_profile is not None + ): # overwrite impurity profile with non-Ne dependent profile + self.impurity_density.loc[ + dict(t=t, element=impurity_to_profile) + ] = self.Nimp_prof() def update_profiles( - self, - parameters: dict, + self, + parameters: dict, ): """ Update plasma profiles with profile parameters i.e. @@ -501,7 +509,7 @@ def update_profiles( if any(imp_conc): for idx, key in enumerate(self.impurities): if any([imp for imp in imp_conc if key in imp]): - self.impurity_concentration[idx] = parameters[key+"_conc"] + self.impurity_concentration[idx] = parameters[key + "_conc"] if any([key for key in parameters.keys() if "Nimp_prof" in key]): impurity_to_profile = "ar" # TODO: generalise this for n independent impurity profiles e.g. Nimp1/Nimp2/... @@ -547,7 +555,7 @@ def pressure_tot(self): def pressure_fast(self): # TODO: check whether degrees of freedom are correctly included... self._pressure_fast.values = ( - self.pressure_fast_parallel / 3 + self.pressure_fast_perpendicular * 2 / 3 + self.pressure_fast_parallel / 3 + self.pressure_fast_perpendicular * 2 / 3 ) return self._pressure_fast @@ -614,8 +622,8 @@ def calc_zeff(self): meanz = self.meanz for elem in self.elements: self._zeff.loc[dict(element=elem)] = ( - (ion_density.sel(element=elem) * meanz.sel(element=elem) ** 2) - / electron_density + (ion_density.sel(element=elem) * meanz.sel(element=elem) ** 2) + / electron_density ).values return self._zeff @@ -659,8 +667,8 @@ def calc_lz_tot(self): self.power_loss_tot[elem]( Te, Fz, Ne=Ne, Nh=Nh, full_run=self.full_run ) - .transpose() - .values + .transpose() + .values ) return self._lz_tot @@ -687,8 +695,8 @@ def calc_lz_sxr(self): self.power_loss_sxr[elem]( Te, Fz, Ne=Ne, Nh=Nh, full_run=self.full_run ) - .transpose() - .values + .transpose() + .values ) return self._lz_sxr @@ -701,9 +709,9 @@ def calc_total_radiation(self): ion_density = self.ion_density for elem in self.elements: total_radiation = ( - lz_tot[elem].sum("ion_charges") - * self.electron_density - * ion_density.sel(element=elem) + lz_tot[elem].sum("ion_charges") + * self.electron_density + * ion_density.sel(element=elem) ) self._total_radiation.loc[dict(element=elem)] = xr.where( total_radiation >= 0, @@ -724,9 +732,9 @@ def calc_sxr_radiation(self): ion_density = self.ion_density for elem in self.elements: sxr_radiation = ( - lz_sxr[elem].sum("ion_charges") - * self.electron_density - * ion_density.sel(element=elem) + lz_sxr[elem].sum("ion_charges") + * self.electron_density + * ion_density.sel(element=elem) ) self._sxr_radiation.loc[dict(element=elem)] = xr.where( sxr_radiation >= 0, @@ -846,7 +854,11 @@ def calc_impurity_density(self, elements): """ for elem in elements: Nimp = self.electron_density * self.impurity_concentration.sel(element=elem) - self.impurity_density.loc[dict(element=elem, )] = Nimp.values + self.impurity_density.loc[ + dict( + element=elem, + ) + ] = Nimp.values def impose_flat_zeff(self): """ @@ -856,14 +868,14 @@ def impose_flat_zeff(self): for elem in self.impurities: if np.count_nonzero(self.ion_density.sel(element=elem)) != 0: zeff_tmp = ( - self.ion_density.sel(element=elem) - * self.meanz.sel(element=elem) ** 2 - / self.electron_density + self.ion_density.sel(element=elem) + * self.meanz.sel(element=elem) ** 2 + / self.electron_density ) value = zeff_tmp.where(zeff_tmp.rho_poloidal < 0.2).mean("rho_poloidal") zeff_tmp = zeff_tmp / zeff_tmp * value ion_density_tmp = zeff_tmp / ( - self.meanz.sel(element=elem) ** 2 / self.electron_density + self.meanz.sel(element=elem) ** 2 / self.electron_density ) self.ion_density.loc[dict(element=elem)] = ion_density_tmp.values @@ -875,13 +887,13 @@ def convert_in_time(self, value: DataArray, method="linear"): return binned def build_atomic_data( - self, - Te: DataArray = None, - Ne: DataArray = None, - Nh: DataArray = None, - tau: DataArray = None, - default=True, - calc_power_loss=True, + self, + Te: DataArray = None, + Ne: DataArray = None, + Nh: DataArray = None, + tau: DataArray = None, + default=True, + calc_power_loss=True, ): if default: xend = 1.02 @@ -925,8 +937,8 @@ def build_atomic_data( power_loss_tot[elem] = PowerLoss(plt, prb, PRC=prc) power_loss_tot[elem](Te, F_z_t, Ne=Ne, Nh=Nh, full_run=self.full_run) if ( - "pls" in self.adf11[elem].keys() - and "prs" in self.adf11[elem].keys() + "pls" in self.adf11[elem].keys() + and "prs" in self.adf11[elem].keys() ): try: pls = self.ADASReader.get_adf11( @@ -992,13 +1004,13 @@ def map_to_midplane(self): for t in np.array(self.time_to_calculate, ndmin=1): rho = ( self.equilibrium.rho.sel(t=t, method="nearest") - .interp(R=R, z=z) - .drop_vars(["R", "z"]) + .interp(R=R, z=z) + .drop_vars(["R", "z"]) ) midplane_profiles[k].append( prof_rho.sel(t=t, method="nearest") - .interp(rho_poloidal=rho) - .drop_vars("rho_poloidal") + .interp(rho_poloidal=rho) + .drop_vars("rho_poloidal") ) midplane_profiles[k] = xr.concat(midplane_profiles[k], "t").assign_coords( t=self.t @@ -1010,7 +1022,7 @@ def map_to_midplane(self): self.midplane_profiles = midplane_profiles def calc_centrifugal_asymmetry( - self, time=None, test_toroidal_rotation=None, plot=False + self, time=None, test_toroidal_rotation=None, plot=False ): """ Calculate (R, z) maps of the ion densities caused by centrifugal asymmetry @@ -1069,12 +1081,12 @@ def calc_centrifugal_asymmetry( self.centrifugal_asymmetry.loc[dict(element=elem)] = asymm asymmetry_factor = asymm.interp(rho_poloidal=self.rho_2d) self.asymmetry_multiplier.loc[dict(element=elem)] = np.exp( - asymmetry_factor * (self.rho_2d.R ** 2 - R_0 ** 2) + asymmetry_factor * (self.rho_2d.R**2 - R_0**2) ) self.ion_density_2d = ( - ion_density.interp(rho_poloidal=self.rho_2d).drop_vars("rho_poloidal") - * self.asymmetry_multiplier + ion_density.interp(rho_poloidal=self.rho_2d).drop_vars("rho_poloidal") + * self.asymmetry_multiplier ) assign_datatype(self.ion_density_2d, ("density", "ion"), "m^-3") @@ -1115,9 +1127,9 @@ def calc_rad_power_2d(self): """ for elem in self.elements: total_radiation = ( - self.lz_tot[elem].sum("ion_charges") - * self.electron_density - * self.ion_density.sel(element=elem) + self.lz_tot[elem].sum("ion_charges") + * self.electron_density + * self.ion_density.sel(element=elem) ) total_radiation = xr.where( total_radiation >= 0, @@ -1127,9 +1139,9 @@ def calc_rad_power_2d(self): self.total_radiation.loc[dict(element=elem)] = total_radiation.values sxr_radiation = ( - self.lz_sxr[elem].sum("ion_charges") - * self.electron_density - * self.ion_density.sel(element=elem) + self.lz_sxr[elem].sum("ion_charges") + * self.electron_density + * self.ion_density.sel(element=elem) ) sxr_radiation = xr.where( sxr_radiation >= 0, @@ -1166,9 +1178,9 @@ def write_to_pickle(self): # Generalized dependency caching class TrackDependecies: def __init__( - self, - operator: Callable, - dependencies: list, + self, + operator: Callable, + dependencies: list, ): """ Call operator only if dependencies variables have changed. @@ -1185,8 +1197,8 @@ def __init__( self.dependencies = dependencies def numpyhash( - self, - nparray: np.array, + self, + nparray: np.array, ): a = nparray.view(np.uint8) return hashlib.sha1(a).hexdigest() @@ -1228,18 +1240,18 @@ def __call__(self): def example_run( - pulse: int = None, - tstart=0.02, - tend=0.1, - dt=0.01, - main_ion="h", - impurities=("c", "ar", "he"), - impurity_concentration=(0.03, 0.001, 0.01), - verbose: bool = True, - n_rad: int = 41, - full_run=False, - calc_power_loss: bool = False, - **kwargs, + pulse: int = None, + tstart=0.02, + tend=0.1, + dt=0.01, + main_ion="h", + impurities=("c", "ar", "he"), + impurity_concentration=(0.03, 0.001, 0.01), + verbose: bool = True, + n_rad: int = 41, + full_run=False, + calc_power_loss: bool = False, + **kwargs, ): # TODO: swap all profiles to new version! diff --git a/indica/workflows/bayes_dev_workflow.py b/indica/workflows/bayes_dev_workflow.py index ac0110d0..c673d2e7 100644 --- a/indica/workflows/bayes_dev_workflow.py +++ b/indica/workflows/bayes_dev_workflow.py @@ -25,19 +25,16 @@ "Ne_prof.y1": 2e18, "Ne_prof.yend": 1e18, "Ne_prof.wped": 2, - "Nimp_prof.y0": 3e16, "Nimp_prof.y1": 0.5e16, "Nimp_prof.wcenter": 0.4, "Nimp_prof.wped": 6, "Nimp_prof.peaking": 2, - "Te_prof.y0": 3000, "Te_prof.y1": 50, "Te_prof.wcenter": 0.4, "Te_prof.wped": 3, "Te_prof.peaking": 2, - "Ti_prof.y0": 6000, "Ti_prof.y1": 50, "Ti_prof.wcenter": 0.4, @@ -52,24 +49,24 @@ "Ne_prof.wped": get_uniform(1, 6), "Ne_prof.wcenter": get_uniform(0.1, 0.8), "Ne_prof.peaking": get_uniform(1, 6), - "ar_conc": loguniform(0.0001, 0.01), - "Nimp_prof.y0": get_uniform(1e16, 1e18), "Nimp_prof.y1": get_uniform(1e15, 2e16), - "Ne_prof.y0/Nimp_prof.y0": lambda x1, x2: np.where((x1 > x2 * 100) & (x1 < x2 * 1e5), 1, 0), + "Ne_prof.y0/Nimp_prof.y0": lambda x1, x2: np.where( + (x1 > x2 * 100) & (x1 < x2 * 1e5), 1, 0 + ), "Nimp_prof.y0/Nimp_prof.y1": lambda x1, x2: np.where((x1 > x2), 1, 0), "Nimp_prof.wped": get_uniform(1, 6), "Nimp_prof.wcenter": get_uniform(0.1, 0.8), "Nimp_prof.peaking": get_uniform(1, 20), - "Nimp_prof.peaking/Ne_prof.peaking": lambda x1, x2: np.where((x1 > x2), 1, 0), # impurity always more peaked - + "Nimp_prof.peaking/Ne_prof.peaking": lambda x1, x2: np.where( + (x1 > x2), 1, 0 + ), # impurity always more peaked "Te_prof.y0": get_uniform(1000, 5000), "Te_prof.wped": get_uniform(1, 6), "Te_prof.wcenter": get_uniform(0.1, 0.6), "Te_prof.peaking": get_uniform(1, 6), "Ti_prof.y0/Te_prof.y0": lambda x1, x2: np.where(x1 > x2, 1, 0), # hot ion mode - "Ti_prof.y0": get_uniform(3000, 10000), "Ti_prof.wped": get_uniform(1, 6), "Ti_prof.wcenter": get_uniform(0.1, 0.6), @@ -82,51 +79,43 @@ "Ne_prof.peaking", "Ne_prof.wcenter", "Ne_prof.wped", - # "ar_conc", # "Nimp_prof.y1", "Nimp_prof.y0", "Nimp_prof.wcenter", "Nimp_prof.wped", "Nimp_prof.peaking", - "Te_prof.y0", "Te_prof.wped", "Te_prof.wcenter", "Te_prof.peaking", - "Ti_prof.y0", "Ti_prof.wped", "Ti_prof.wcenter", "Ti_prof.peaking", ] -OPTIMISED_QUANTITY = [ - "xrcs.spectra", - "cxff_pi.ti", - "efit.wp", - "smmh1.ne" -] +OPTIMISED_QUANTITY = ["xrcs.spectra", "cxff_pi.ti", "efit.wp", "smmh1.ne"] class BayesWorkflow: - def __init__(self, - pulse=None, - result_path="./results/example/", - nwalkers=50, - tstart=0.02, - tend=0.10, - dt=0.01, - tsample=0.06, - iterations=100, - burn_in_fraction=0, - center_mass_sampling=True, - Ti_ref=True, - diagnostics=None, - profiles=None, - phantom=True, - ): - + def __init__( + self, + pulse=None, + result_path="./results/example/", + nwalkers=50, + tstart=0.02, + tend=0.10, + dt=0.01, + tsample=0.06, + iterations=100, + burn_in_fraction=0, + center_mass_sampling=True, + Ti_ref=True, + diagnostics=None, + profiles=None, + phantom=True, + ): self.pulse = pulse self.tstart = tstart self.tend = tend @@ -149,13 +138,18 @@ def __init__(self, dt=dt, main_ion="h", impurities=("ar", "c"), - impurity_concentration=(0.001, 0.04,), + impurity_concentration=( + 0.001, + 0.04, + ), full_run=False, n_rad=20, ) self.plasma.Ti_ref = self.Ti_ref self.tsample = tsample - self.plasma.time_to_calculate = self.plasma.t[np.abs(tsample-self.plasma.t).argmin()] + self.plasma.time_to_calculate = self.plasma.t[ + np.abs(tsample - self.plasma.t).argmin() + ] self.plasma.update_profiles(DEFAULT_PHANTOM_PARAMS) self.plasma.build_atomic_data(calc_power_loss=False) @@ -185,16 +179,24 @@ def __init__(self, log_prob_fn=self.bayes_run.ln_posterior, parameter_names=OPTIMISED_PARAMS, moves=self.move, - kwargs={"moment_analysis": False, "calc_spectra": True, "minimum_lines": False, - "background": self.flat_data["xrcs.background"]}, + kwargs={ + "moment_analysis": False, + "calc_spectra": True, + "minimum_lines": False, + "background": self.flat_data["xrcs.background"], + }, ) if not self.center_mass_sampling: - self.start_points = self.bayes_run.sample_from_priors(OPTIMISED_PARAMS, size=self.nwalkers) + self.start_points = self.bayes_run.sample_from_priors( + OPTIMISED_PARAMS, size=self.nwalkers + ) else: - # gaussian mass around most probable starting points + # gaussian mass around most probable starting points TODO: move this function to it's own method nsamples = 100 - start_points = self.bayes_run.sample_from_priors(OPTIMISED_PARAMS, size=nsamples) + start_points = self.bayes_run.sample_from_priors( + OPTIMISED_PARAMS, size=nsamples + ) ln_prob, _ = self.sampler.compute_log_prob(start_points) num_best_points = int(nsamples * 0.05) index_best_start = np.argsort(ln_prob)[-num_best_points:] @@ -204,26 +206,37 @@ def __init__(self, # Passing samples through ln_prior and redrawing if they fail samples = np.empty((OPTIMISED_PARAMS.__len__(), 0)) while samples.size < OPTIMISED_PARAMS.__len__() * self.nwalkers: - sample = np.random.normal(np.mean(best_start_points, axis=0), best_points_std * 2, - size=(self.nwalkers * 5, len(OPTIMISED_PARAMS))) - start = {name: sample[:, idx] for idx, name in enumerate(OPTIMISED_PARAMS)} + sample = np.random.normal( + np.mean(best_start_points, axis=0), + best_points_std * 2, + size=(self.nwalkers * 5, len(OPTIMISED_PARAMS)), + ) + start = { + name: sample[:, idx] for idx, name in enumerate(OPTIMISED_PARAMS) + } ln_prior = self.bayes_run._ln_prior(start) # Convert from dictionary of arrays -> array, # then filtering out where ln_prior is -infinity accepted_samples = np.array(list(start.values()))[ - :, ln_prior != -np.inf - ] + :, ln_prior != -np.inf + ] samples = np.append(samples, accepted_samples, axis=1) - self.start_points = samples[:, 0:self.nwalkers].transpose() + self.start_points = samples[:, 0 : self.nwalkers].transpose() def init_fast_particles(self): st40_code = ReadST40(13110009, self.tstart, self.tend, dt=self.dt, tree="astra") st40_code.get_raw_data("", "astra", "RUN573") st40_code.bin_data_in_time(["astra"], self.tstart, self.tend, self.dt) code_data = st40_code.binned_data["astra"] - Nf = code_data["nf"].interp(rho_poloidal=self.plasma.rho, t=self.plasma.t) * 1.0e19 + Nf = ( + code_data["nf"].interp(rho_poloidal=self.plasma.rho, t=self.plasma.t) + * 1.0e19 + ) self.plasma.fast_density.values = Nf.values - Nn = code_data["nn"].interp(rho_poloidal=self.plasma.rho, t=self.plasma.t) * 1.0e19 + Nn = ( + code_data["nn"].interp(rho_poloidal=self.plasma.rho, t=self.plasma.t) + * 1.0e19 + ) self.plasma.neutral_density.values = Nn.values Pblon = code_data["pblon"].interp(rho_poloidal=self.plasma.rho, t=self.plasma.t) self.plasma.pressure_fast_parallel.values = Pblon.values @@ -231,7 +244,9 @@ def init_fast_particles(self): self.plasma.pressure_fast_perpendicular.values = Pbper.values def read_st40(self, diagnostics=None): - self.ST40_data = ReadST40(self.pulse, tstart=self.tstart, tend=self.tend, dt=self.dt) + self.ST40_data = ReadST40( + self.pulse, tstart=self.tstart, tend=self.tend, dt=self.dt + ) self.ST40_data(diagnostics) self.plasma.set_equilibrium(self.ST40_data.equilibrium) @@ -276,7 +291,8 @@ def init_models(self): direction[:, 2], name="", machine_dimensions=machine_dims, - passes=2, ) + passes=2, + ) los_transform.set_equilibrium(self.plasma.equilibrium) model = Interferometry(name="smmh1") model.set_los_transform(los_transform) @@ -285,9 +301,14 @@ def init_models(self): if "xrcs" in self.diagnostics: los_transform = self.ST40_data.binned_data["xrcs"]["te_kw"].transform - model = Helike_spectroscopy(name="xrcs", window_masks=[slice(0.394, 0.396)], - window_vector=self.ST40_data.binned_data["xrcs"][ - "spectra"].wavelength.values * 0.1) + model = Helike_spectroscopy( + name="xrcs", + window_masks=[slice(0.394, 0.396)], + window_vector=self.ST40_data.binned_data["xrcs"][ + "spectra" + ].wavelength.values + * 0.1, + ) model.set_los_transform(los_transform) model.plasma = self.plasma self.models["xrcs"] = model @@ -308,51 +329,80 @@ def init_models(self): def phantom_data(self, noise=False, noise_factor=0.1): self.flat_data = {} if "smmh1" in self.diagnostics: - self.flat_data["smmh1.ne"] = self.models["smmh1"]().pop("ne").expand_dims( - dim={"t": [self.plasma.time_to_calculate]}) + self.flat_data["smmh1.ne"] = ( + self.models["smmh1"]() + .pop("ne") + .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) + ) if "xrcs" in self.diagnostics: - self.flat_data["xrcs.spectra"] = self.models["xrcs"]().pop("spectra").expand_dims( - dim={"t": [self.plasma.time_to_calculate]}) + self.flat_data["xrcs.spectra"] = ( + self.models["xrcs"]() + .pop("spectra") + .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) + ) self.flat_data["xrcs.background"] = None if "cxff_pi" in self.diagnostics: - cxrs_data = self.models["cxff_pi"]().pop("ti").expand_dims(dim={"t": [self.plasma.time_to_calculate]}) + cxrs_data = ( + self.models["cxff_pi"]() + .pop("ti") + .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) + ) self.flat_data["cxff_pi.ti"] = cxrs_data.where(cxrs_data.channel == 2) if "efit" in self.diagnostics: - self.flat_data["efit.wp"] = self.models["efit"]().pop("wp").expand_dims( - dim={"t": [self.plasma.time_to_calculate]}) + self.flat_data["efit.wp"] = ( + self.models["efit"]() + .pop("wp") + .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) + ) if noise: self.flat_data["smmh1.ne"] = self.flat_data["smmh1.ne"] + self.flat_data[ - "smmh1.ne"].max().values * np.random.normal(0, noise_factor, None) - self.flat_data["xrcs.spectra"] = self.flat_data["xrcs.spectra"] + np.random.normal(0, np.sqrt( - self.flat_data["xrcs.spectra"].values[0,]), self.flat_data["xrcs.spectra"].shape[1]) - self.flat_data["cxff_pi.ti"] = self.flat_data["cxff_pi.ti"] + self.flat_data[ - "cxff_pi.ti"].max().values * np.random.normal(0, noise_factor, self.flat_data["cxff_pi.ti"].shape[1]) + "smmh1.ne" + ].max().values * np.random.normal(0, noise_factor, None) + self.flat_data["xrcs.spectra"] = self.flat_data[ + "xrcs.spectra" + ] + np.random.normal( + 0, + np.sqrt(self.flat_data["xrcs.spectra"].values[0,]), + self.flat_data["xrcs.spectra"].shape[1], + ) + self.flat_data["cxff_pi.ti"] = self.flat_data[ + "cxff_pi.ti" + ] + self.flat_data["cxff_pi.ti"].max().values * np.random.normal( + 0, noise_factor, self.flat_data["cxff_pi.ti"].shape[1] + ) self.flat_data["efit.wp"] = self.flat_data["efit.wp"] + self.flat_data[ - "efit.wp"].max().values * np.random.normal(0, noise_factor, None) + "efit.wp" + ].max().values * np.random.normal(0, noise_factor, None) def exp_data(self): self.flat_data = flatdict.FlatDict(self.ST40_data.binned_data, ".") if "xrcs" in self.diagnostics: - self.flat_data["xrcs.spectra"]["wavelength"] = self.flat_data["xrcs.spectra"].wavelength * 0.1 + self.flat_data["xrcs.spectra"]["wavelength"] = ( + self.flat_data["xrcs.spectra"].wavelength * 0.1 + ) background = self.flat_data["xrcs.spectra"].where( - (self.flat_data["xrcs.spectra"].wavelength < 0.392) & - (self.flat_data["xrcs.spectra"].wavelength > 0.388), - drop=True) + (self.flat_data["xrcs.spectra"].wavelength < 0.392) + & (self.flat_data["xrcs.spectra"].wavelength > 0.388), + drop=True, + ) self.flat_data["xrcs.background"] = background.mean(dim="wavelength") self.flat_data["xrcs.spectra"]["error"] = np.sqrt( - self.flat_data["xrcs.spectra"] + background.std(dim="wavelength") ** 2) + self.flat_data["xrcs.spectra"] + background.std(dim="wavelength") ** 2 + ) if "cxff_pi" in self.diagnostics: self.flat_data["cxff_pi"]["ti"] = self.flat_data["cxff_pi"]["ti"].where( - self.flat_data["cxff_pi"]["ti"].channel == 2, drop=True) + self.flat_data["cxff_pi"]["ti"].channel == 2, drop=True + ) def __call__(self, *args, **kwargs): - autocorr = sample_with_autocorr( self.sampler, self.start_points, iterations=self.iterations, auto_sample=5 ) - blobs = self.sampler.get_blobs(discard=int(self.iterations * self.burn_in_fraction), flat=True) + blobs = self.sampler.get_blobs( + discard=int(self.iterations * self.burn_in_fraction), flat=True + ) blob_names = self.sampler.get_blobs().flatten()[0].keys() blob_dict = { blob_name: xr.concat( @@ -363,7 +413,9 @@ def __call__(self, *args, **kwargs): } samples = self.sampler.get_chain(flat=True) - prior_samples = self.bayes_run.sample_from_priors(OPTIMISED_PARAMS, size=int(1e4)) + prior_samples = self.bayes_run.sample_from_priors( + OPTIMISED_PARAMS, size=int(1e4) + ) result = { "blobs": blob_dict, "diag_data": self.flat_data, @@ -379,12 +431,17 @@ def __call__(self, *args, **kwargs): if __name__ == "__main__": - run = BayesWorkflow(pulse=10009, result_path="./results/10009_60ms_long/", iterations=500, nwalkers=200, - burn_in_fraction=0.10, dt=0.005, tsample=0.060, - diagnostics=[ - "xrcs", - "efit", - "smmh1", - "cxff_pi" - ], phantom=False, center_mass_sampling=True, Ti_ref=True) + run = BayesWorkflow( + pulse=10009, + result_path="./results/10009_60ms_long/", + iterations=500, + nwalkers=200, + burn_in_fraction=0.10, + dt=0.005, + tsample=0.060, + diagnostics=["xrcs", "efit", "smmh1", "cxff_pi"], + phantom=False, + center_mass_sampling=True, + Ti_ref=True, + ) run() diff --git a/indica/workflows/bayes_workflow.py b/indica/workflows/bayes_workflow.py index e29add98..781bce8d 100644 --- a/indica/workflows/bayes_workflow.py +++ b/indica/workflows/bayes_workflow.py @@ -9,16 +9,16 @@ def plot_profile( - profile, - blobkey: str, - figheader="./results/test/", - phantom_profile=None, - logscale=False, - sharefig=False, - filename="", - filetype=".png", - linestyle="--", - color="blue", + profile, + blobkey: str, + figheader="./results/test/", + phantom_profile=None, + logscale=False, + sharefig=False, + filename="", + filetype=".png", + linestyle="--", + color="blue", ): set_plot_rcparams("profiles") # if not plt.get_fignums(): # If no figure is open @@ -96,16 +96,16 @@ def plot_profile( def _plot_1d( - blobs: dict, - blobkey: str, - diag_data: dict, - filename: str, - figheader="./results/test/", - ylabel="a.u.", - xlabel="[]", - xlim=None, - figsize=(6.4, 4.8), - **kwargs, + blobs: dict, + blobkey: str, + diag_data: dict, + filename: str, + figheader="./results/test/", + ylabel="a.u.", + xlabel="[]", + xlim=None, + figsize=(6.4, 4.8), + **kwargs, ): if blobkey not in blobs.keys(): raise ValueError(f"{blobkey} not in blobs") @@ -162,24 +162,39 @@ def _plot_1d( plt.close() -def violinplot(data: dict, key: str, diag_data: dict, filename: str, ylabel="[a.u.]", figheader="./results/test/", - **kwargs): +def violinplot( + data: dict, + key: str, + diag_data: dict, + filename: str, + ylabel="[a.u.]", + figheader="./results/test/", + **kwargs, +): set_plot_rcparams() - fig, axs = plt.subplots(1, 1,) + fig, axs = plt.subplots( + 1, + 1, + ) _data = data[key].values - _data = _data[((_data > np.quantile(_data, 0.16)) & (_data < np.quantile(_data, 0.84)))] - violin = axs.violinplot(_data, - showextrema=False, - # quantiles=[0.025, 0.975, 0.16, 0.84], - # showmedians=True, - ) + _data = _data[ + ((_data > np.quantile(_data, 0.16)) & (_data < np.quantile(_data, 0.84))) + ] + violin = axs.violinplot( + _data, + showextrema=False, + # quantiles=[0.025, 0.975, 0.16, 0.84], + # showmedians=True, + ) violin["bodies"][0].set_edgecolor("black") axs.set_xlabel(key) top = axs.get_ylim()[1] axs.set_ylim(bottom=0, top=top * 1.1) axs.set_ylabel(f"{ylabel}") y = diag_data[key].sel(t=data[key].t).values - axs.errorbar(1, y=y, yerr=y * 0.10, fmt="D", ecolor="black", capsize=10, color="black") + axs.errorbar( + 1, y=y, yerr=y * 0.10, fmt="D", ecolor="black", capsize=10, color="black" + ) set_axis_sci() plt.setp([axs.get_xticklabels()], visible=False) plt.savefig(figheader + filename) @@ -217,10 +232,10 @@ def histograms(data, diag_data, filename): def plot_bayes_result( - results, - figheader="./results/test/", - filetype=".png", - **kwargs, + results, + figheader="./results/test/", + filetype=".png", + **kwargs, ): diag_data = results["diag_data"] blobs = results["blobs"] @@ -247,8 +262,12 @@ def plot_bayes_result( plt.close() if "cxff_pi.ti" in blobs.keys(): - blobs["cxff_pi.ti0"] = blobs["cxff_pi.ti"].sel(channel=diag_data["cxff_pi.ti"].channel) - diag_data["cxff_pi.ti0"] = diag_data["cxff_pi.ti"].sel(channel=diag_data["cxff_pi.ti"].channel) + blobs["cxff_pi.ti0"] = blobs["cxff_pi.ti"].sel( + channel=diag_data["cxff_pi.ti"].channel + ) + diag_data["cxff_pi.ti0"] = diag_data["cxff_pi.ti"].sel( + channel=diag_data["cxff_pi.ti"].channel + ) key = "cxff_pi.ti0" if key in blobs.keys(): @@ -324,7 +343,6 @@ def plot_bayes_result( figheader=figheader, ylabel="Temperature [eV]", xlabel="Channel", - ) key = "electron_temperature" @@ -351,7 +369,11 @@ def plot_bayes_result( ) key = "electron_density" plot_profile( - blobs[key], key, figheader=figheader, filetype=filetype, phantom_profile=phantom_profiles + blobs[key], + key, + figheader=figheader, + filetype=filetype, + phantom_profile=phantom_profiles, ) key = "impurity_density" for elem in blobs[key].element.values: @@ -391,7 +413,7 @@ def plot_bayes_result( figheader=figheader, filetype=filetype, phantom_profile=phantom_profiles, - logscale=True + logscale=True, ) corner.corner(samples, labels=param_names) @@ -406,9 +428,9 @@ def sample_with_autocorr(sampler, start_points, iterations=10, auto_sample=5): autocorr = np.ones((iterations,)) * np.nan old_tau = np.inf for sample in sampler.sample( - start_points, - iterations=iterations, - progress=True, + start_points, + iterations=iterations, + progress=True, ): if sampler.iteration % auto_sample: continue From dfabdd0bc5eb6dabee45e711365b987b977e9bb1 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Wed, 31 May 2023 14:05:54 +0100 Subject: [PATCH 049/160] deleted example bayes_opt --- indica/workflows/example_bayes_opt.py | 180 -------------------------- 1 file changed, 180 deletions(-) delete mode 100644 indica/workflows/example_bayes_opt.py diff --git a/indica/workflows/example_bayes_opt.py b/indica/workflows/example_bayes_opt.py deleted file mode 100644 index 4ac1bcee..00000000 --- a/indica/workflows/example_bayes_opt.py +++ /dev/null @@ -1,180 +0,0 @@ -import emcee -import numpy as np -import pandas as pd -import xarray as xr - -from indica.bayesmodels import BayesModels -from indica.bayesmodels import get_uniform -from indica.models.helike_spectroscopy import Helike_spectroscopy -from indica.models.interferometry import Interferometry -from indica.models.plasma import Plasma -from indica.readers.read_st40 import ReadST40 -from indica.workflows.bayes_workflow import plot_bayes_result -from indica.workflows.bayes_workflow import sample_with_autocorr - -# TODO: allow conditional prior usage even when only -# one param is being optimisied i.e. 1 is constant - - -def run( - pulse, - phantom_profile_params, - iterations, - result_path, - burn_in=0, - tstart=0.02, - tend=0.10, - dt=0.01, - tsample=3, -): - plasma = Plasma( - tstart=tstart, - tend=tend, - dt=dt, - main_ion="h", - impurities=("ar",), - impurity_concentration=(0.001,), - full_run=False, - n_rad=10, - ) - plasma.time_to_calculate = plasma.t[tsample] - plasma.update_profiles(phantom_profile_params) - plasma.build_atomic_data() - # Make phantom profiles - phantom_profiles = { - "electron_density": plasma.Ne_prof.yspl, - "electron_temperature": plasma.Te_prof.yspl, - "ion_temperature": plasma.Ti_prof.yspl, - "impurity_density": plasma.Nimp_prof.yspl, - } - - ST40 = ReadST40(pulse, tstart=tstart, tend=tend) - ST40(["xrcs", "smmh1"]) - - # Initialise Diagnostic Models - los_transform = ST40.binned_data["smmh1"]["ne"].transform - smmh1 = Interferometry(name="smmh1") - smmh1.set_los_transform(los_transform) - smmh1.plasma = plasma - los_transform = ST40.binned_data["xrcs"]["te_kw"].transform - xrcs = Helike_spectroscopy(name="xrcs", window_masks=[slice(0.3945, 0.3962)]) - xrcs.set_los_transform(los_transform) - xrcs.plasma = plasma - - flat_data = {} - flat_data["smmh1.ne"] = ( - smmh1().pop("ne").expand_dims(dim={"t": [plasma.time_to_calculate]}) - ) - flat_data["xrcs.spectra"] = ( - xrcs().pop("spectra").expand_dims(dim={"t": [plasma.time_to_calculate]}) - ) - - priors = { - "Ne_prof.y0": get_uniform(1e19, 8e19), - "Ne_prof.y1": get_uniform(1e18, 5e18), - "Ne_prof.y0/Ne_prof.y1": lambda x1, x2: np.where(((x1 > x2 * 2)), 1, 0), - "Ne_prof.wped": get_uniform(1, 5), - "Ne_prof.wcenter": get_uniform(0.1, 0.8), - "Ne_prof.peaking": get_uniform(1, 5), - "Nimp_prof.peaking": get_uniform(1, 8), - "Nimp_prof.wcenter": get_uniform(0.1, 0.4), - "Nimp_prof.y0": get_uniform(1e16, 5e16), - "Nimp_prof.y1": get_uniform(1e15, 1e16), - "Ne_prof.y0/Nimp_prof.y0": lambda x1, x2: np.where( - (x1 > x2 * 100) & (x1 < x2 * 1e4), 1, 0 - ), - "Nimp_prof.y0/Nimp_prof.y1": lambda x1, x2: np.where((x1 > x2), 1, 0), - "Te_prof.y0": get_uniform(1000, 6000), - "Te_prof.peaking": get_uniform(1, 4), - "Ti_prof.y0": get_uniform(2000, 10000), - "Ti_prof.peaking": get_uniform(1, 4), - } - # Setup Optimiser - param_names = [ - "Ne_prof.y0", - # "Ne_prof.y1", - # "Ne_prof.peaking", - "Nimp_prof.y0", - # "Nimp_prof.y1", - # "Nimp_prof.peaking", - "Te_prof.y0", - # "Te_prof.peaking", - "Ti_prof.y0", - # "Ti_prof.peaking", - ] - - bm = BayesModels( - plasma=plasma, - data=flat_data, - diagnostic_models=[smmh1, xrcs], - quant_to_optimise=[ - "smmh1.ne", - "xrcs.spectra", - ], - priors=priors, - ) - - ndim = param_names.__len__() - nwalkers = 20 - start_points = bm.sample_from_priors(param_names, size=nwalkers) - move = [(emcee.moves.StretchMove(), 1.0), (emcee.moves.DEMove(), 0.0)] - - sampler = emcee.EnsembleSampler( - nwalkers, - ndim, - log_prob_fn=bm.ln_posterior, - parameter_names=param_names, - moves=move, - kwargs={"moment_analysis": False, "calc_spectra": True}, - ) - - autocorr = sample_with_autocorr( - sampler, start_points, iterations=iterations, auto_sample=5 - ) - - blobs = sampler.get_blobs(discard=burn_in, flat=True) - blob_names = sampler.get_blobs().flatten()[0].keys() - blob_dict = { - blob_name: xr.concat( - [data[blob_name] for data in blobs], - dim=pd.Index(np.arange(0, blobs.__len__()), name="index"), - ) - for blob_name in blob_names - } - samples = sampler.get_chain(flat=True) - - prior_samples = bm.sample_from_priors(param_names, size=int(1e5)) - - # TODO make sure xrcs bckc doesn't have dims t and channels - # save result - result = { - "blobs": blob_dict, - "diag_data": flat_data, - "samples": samples, - "prior_samples": prior_samples, - "param_names": param_names, - "phantom_profiles": phantom_profiles, - "plasma": plasma, - "autocorr": autocorr, - } - print(sampler.acceptance_fraction.sum()) - plot_bayes_result(**result, figheader=result_path) - - -if __name__ == "__main__": - params = { - "Ne_prof.y0": 5e19, - "Ne_prof.wcenter": 0.4, - "Ne_prof.peaking": 2, - "Ne_prof.y1": 2e18, - "Ne_prof.yend": 1e18, - "Ne_prof.wped": 2, - "Nimp_prof.y0": 2e16, - "Nimp_prof.y1": 2e15, - "Nimp_prof.peaking": 2, - "Te_prof.y0": 3000, - "Te_prof.peaking": 2, - "Ti_prof.y0": 5000, - "Ti_prof.peaking": 2, - } - run(10009, params, 10, "./results/test/", burn_in=0) From b21b66d2007d793b8da970b83e955f391d2e164a Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Wed, 31 May 2023 14:07:40 +0100 Subject: [PATCH 050/160] moved sample with autocorr function to bayes_dev_workflow --- indica/workflows/bayes_dev_workflow.py | 23 ++++++++++++++++++++++- indica/workflows/bayes_workflow.py | 21 --------------------- 2 files changed, 22 insertions(+), 22 deletions(-) diff --git a/indica/workflows/bayes_dev_workflow.py b/indica/workflows/bayes_dev_workflow.py index c673d2e7..a0e4eaba 100644 --- a/indica/workflows/bayes_dev_workflow.py +++ b/indica/workflows/bayes_dev_workflow.py @@ -8,7 +8,7 @@ from indica.readers.read_st40 import ReadST40 from indica.bayesmodels import BayesModels, get_uniform -from indica.workflows.bayes_workflow import plot_bayes_result, sample_with_autocorr +from indica.workflows.bayes_workflow import plot_bayes_result from indica.models.interferometry import Interferometry from indica.models.helike_spectroscopy import Helike_spectroscopy @@ -430,6 +430,27 @@ def __call__(self, *args, **kwargs): plot_bayes_result(result, figheader=self.result_path, filetype=".png") +def sample_with_autocorr(sampler, start_points, iterations=10, auto_sample=5): + autocorr = np.ones((iterations,)) * np.nan + old_tau = np.inf + for sample in sampler.sample( + start_points, + iterations=iterations, + progress=True, + ): + if sampler.iteration % auto_sample: + continue + new_tau = sampler.get_autocorr_time(tol=0) + autocorr[sampler.iteration - 1] = np.mean(new_tau) + converged = np.all(new_tau * 50 < sampler.iteration) + converged &= np.all(np.abs(old_tau - new_tau) / new_tau < 0.01) + if converged: + break + old_tau = new_tau + autocorr = autocorr[: sampler.iteration] + return autocorr + + if __name__ == "__main__": run = BayesWorkflow( pulse=10009, diff --git a/indica/workflows/bayes_workflow.py b/indica/workflows/bayes_workflow.py index 781bce8d..69ed7a46 100644 --- a/indica/workflows/bayes_workflow.py +++ b/indica/workflows/bayes_workflow.py @@ -424,27 +424,6 @@ def plot_bayes_result( plt.close("all") -def sample_with_autocorr(sampler, start_points, iterations=10, auto_sample=5): - autocorr = np.ones((iterations,)) * np.nan - old_tau = np.inf - for sample in sampler.sample( - start_points, - iterations=iterations, - progress=True, - ): - if sampler.iteration % auto_sample: - continue - new_tau = sampler.get_autocorr_time(tol=0) - autocorr[sampler.iteration - 1] = np.mean(new_tau) - converged = np.all(new_tau * 50 < sampler.iteration) - converged &= np.all(np.abs(old_tau - new_tau) / new_tau < 0.01) - if converged: - break - old_tau = new_tau - autocorr = autocorr[: sampler.iteration] - return autocorr - - if __name__ == "__main__": filehead = "./results/10009_60ms_short/" with open(filehead + "results.pkl", "rb") as handle: From 2f78d9d745ebab776f42eada4e9ed4da3261fb5b Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Wed, 31 May 2023 14:08:53 +0100 Subject: [PATCH 051/160] renamed bayesworkflow --- indica/workflows/bayes_dev_workflow.py | 2 +- indica/workflows/{bayes_workflow.py => bayes_plots.py} | 0 2 files changed, 1 insertion(+), 1 deletion(-) rename indica/workflows/{bayes_workflow.py => bayes_plots.py} (100%) diff --git a/indica/workflows/bayes_dev_workflow.py b/indica/workflows/bayes_dev_workflow.py index a0e4eaba..f17e1c6e 100644 --- a/indica/workflows/bayes_dev_workflow.py +++ b/indica/workflows/bayes_dev_workflow.py @@ -8,7 +8,7 @@ from indica.readers.read_st40 import ReadST40 from indica.bayesmodels import BayesModels, get_uniform -from indica.workflows.bayes_workflow import plot_bayes_result +from indica.workflows.bayes_plots import plot_bayes_result from indica.models.interferometry import Interferometry from indica.models.helike_spectroscopy import Helike_spectroscopy diff --git a/indica/workflows/bayes_workflow.py b/indica/workflows/bayes_plots.py similarity index 100% rename from indica/workflows/bayes_workflow.py rename to indica/workflows/bayes_plots.py From 9ed488ea0ec750717f5c5d1657fab2df275ce1df Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 13 Jun 2023 09:40:12 +0100 Subject: [PATCH 052/160] stashing --- indica/workflows/bayes_dev_workflow.py | 94 ++++++++++++++++++++------ indica/workflows/bayes_plots.py | 80 +++++++++++----------- 2 files changed, 116 insertions(+), 58 deletions(-) diff --git a/indica/workflows/bayes_dev_workflow.py b/indica/workflows/bayes_dev_workflow.py index f17e1c6e..d6e72780 100644 --- a/indica/workflows/bayes_dev_workflow.py +++ b/indica/workflows/bayes_dev_workflow.py @@ -3,13 +3,13 @@ import xarray as xr import pandas as pd import flatdict - from scipy.stats import loguniform +from pathlib import Path +import pickle from indica.readers.read_st40 import ReadST40 from indica.bayesmodels import BayesModels, get_uniform from indica.workflows.bayes_plots import plot_bayes_result - from indica.models.interferometry import Interferometry from indica.models.helike_spectroscopy import Helike_spectroscopy from indica.models.charge_exchange import ChargeExchange @@ -17,6 +17,8 @@ from indica.models.plasma import Plasma from indica.converters.line_of_sight import LineOfSightTransform +from indica.writers.bda_tree import create_nodes, write_nodes + # global configurations DEFAULT_PHANTOM_PARAMS = { "Ne_prof.y0": 5e19, @@ -82,17 +84,17 @@ # "ar_conc", # "Nimp_prof.y1", "Nimp_prof.y0", - "Nimp_prof.wcenter", - "Nimp_prof.wped", - "Nimp_prof.peaking", + # "Nimp_prof.wcenter", + # "Nimp_prof.wped", + # "Nimp_prof.peaking", "Te_prof.y0", - "Te_prof.wped", - "Te_prof.wcenter", - "Te_prof.peaking", + # "Te_prof.wped", + # "Te_prof.wcenter", + # "Te_prof.peaking", "Ti_prof.y0", - "Ti_prof.wped", - "Ti_prof.wcenter", - "Ti_prof.peaking", + # "Ti_prof.wped", + # "Ti_prof.wcenter", + # "Ti_prof.peaking", ] OPTIMISED_QUANTITY = ["xrcs.spectra", "cxff_pi.ti", "efit.wp", "smmh1.ne"] @@ -102,6 +104,7 @@ class BayesWorkflow: def __init__( self, pulse=None, + pulse_to_write = 23000101, result_path="./results/example/", nwalkers=50, tstart=0.02, @@ -110,13 +113,17 @@ def __init__( tsample=0.06, iterations=100, burn_in_fraction=0, + + diagnostics=None, center_mass_sampling=True, Ti_ref=True, - diagnostics=None, profiles=None, phantom=True, + mds_write=False, + save_plots=True, ): self.pulse = pulse + self.pulse_to_write = pulse_to_write self.tstart = tstart self.tend = tend self.dt = dt @@ -131,6 +138,8 @@ def __init__( self.diagnostics = diagnostics self.phantom = phantom self.profiles = profiles + self.mds_write = mds_write + self.save_plots = save_plots self.plasma = Plasma( tstart=tstart, @@ -416,7 +425,7 @@ def __call__(self, *args, **kwargs): prior_samples = self.bayes_run.sample_from_priors( OPTIMISED_PARAMS, size=int(1e4) ) - result = { + self.results = { "blobs": blob_dict, "diag_data": self.flat_data, "samples": samples, @@ -424,10 +433,44 @@ def __call__(self, *args, **kwargs): "param_names": OPTIMISED_PARAMS, "phantom_profiles": self.profiles, "autocorr": autocorr, + "acceptance_fraction": self.sampler.acceptance_fraction.sum(), + } + + self.nested_results = { + "DIAG_DATA":{"efit":{"wp":1}}, + # self.ST40_data.binned_data["efit"]["wp"].values[4,] + # "MODEL_DATA": {}, + "METADATA": { + "DIAG_RUNS": "EMPTY", + "EQUIL": self.ST40_data.equilibrium.__str__(), + "IMP1": self.plasma.impurities[0], + "MAIN_ION": self.plasma.main_ion, + }, + "OPTIMISATION": { + "AUTO_CORR":autocorr, + "BURN_FRAC":self.burn_in_fraction, + "ITER": self.iterations, + "NWALKERS": self.nwalkers, + "PARAM_NAMES": OPTIMISED_PARAMS, + "POST_SAMPLE": samples, + "PRIOR_SAMPLE": prior_samples + }, + # "PHANTOMS": self.profiles, + # "PROFILES": {}, + # "TIME": self.plasma.t, + # "TIME_OPT": self.plasma.time_to_calculate + } - self.acceptance_fraction = self.sampler.acceptance_fraction.sum() - print(self.acceptance_fraction) - plot_bayes_result(result, figheader=self.result_path, filetype=".png") + + # self.acceptance_fraction = self.sampler.acceptance_fraction.sum() + print(self.results["acceptance_fraction"]) + + Path(self.result_path).mkdir(parents=True, exist_ok=True) + with open(self.result_path + "results.pkl", "wb") as handle: + pickle.dump(self.results, handle) + + if self.save_plots: + plot_bayes_result(self.results, figheader=self.result_path, filetype=".png") def sample_with_autocorr(sampler, start_points, iterations=10, auto_sample=5): @@ -451,12 +494,20 @@ def sample_with_autocorr(sampler, start_points, iterations=10, auto_sample=5): return autocorr + + if __name__ == "__main__": + + mds_write=False + pulse_to_write = 23000101 + if mds_write: + node_structure = create_nodes(pulse_to_write=pulse_to_write, diagnostic_quantities=OPTIMISED_QUANTITY) + run = BayesWorkflow( pulse=10009, - result_path="./results/10009_60ms_long/", - iterations=500, - nwalkers=200, + result_path="./results/test/", + iterations=20, + nwalkers=50, burn_in_fraction=0.10, dt=0.005, tsample=0.060, @@ -464,5 +515,10 @@ def sample_with_autocorr(sampler, start_points, iterations=10, auto_sample=5): phantom=False, center_mass_sampling=True, Ti_ref=True, + mds_write=mds_write, + save_plots=True ) run() + + if mds_write: + write_nodes(pulse_to_write, node_structure, run.nested_results) diff --git a/indica/workflows/bayes_plots.py b/indica/workflows/bayes_plots.py index 69ed7a46..3a76d9a3 100644 --- a/indica/workflows/bayes_plots.py +++ b/indica/workflows/bayes_plots.py @@ -21,8 +21,6 @@ def plot_profile( color="blue", ): set_plot_rcparams("profiles") - # if not plt.get_fignums(): # If no figure is open - # plt.figure(figsize=(8, 6)) if blobkey == "electron_temperature": legkey = "Te" @@ -45,26 +43,27 @@ def plot_profile( label=f"{legkey}, 68% Confidence", zorder=3, color=color, - alpha=0.8, - ) - plt.fill_between( - profile.rho_poloidal, - profile.quantile(0.025, dim="index"), - profile.quantile(0.975, dim="index"), - label=f"{legkey}, 95% Confidence", - zorder=2, - color="grey", - alpha=0.6, - ) - plt.fill_between( - profile.rho_poloidal, - profile.quantile(0.00, dim="index"), - profile.quantile(1.00, dim="index"), - label=f"{legkey}, Max-Min", - zorder=1, - color="lightgrey", - alpha=0.6, + alpha=0.9, ) + if legkey != "Nfast": + plt.fill_between( + profile.rho_poloidal, + profile.quantile(0.025, dim="index"), + profile.quantile(0.975, dim="index"), + label=f"{legkey}, 95% Confidence", + zorder=2, + color="grey", + alpha=0.4, + ) + plt.fill_between( + profile.rho_poloidal, + profile.quantile(0.00, dim="index"), + profile.quantile(1.00, dim="index"), + label=f"{legkey}, Max-Min", + zorder=1, + color="lightgrey", + alpha=0.2, + ) if phantom_profile is not None: if "element" in phantom_profile[blobkey].dims: @@ -237,6 +236,7 @@ def plot_bayes_result( filetype=".png", **kwargs, ): + diag_data = results["diag_data"] blobs = results["blobs"] samples = results["samples"] @@ -245,10 +245,6 @@ def plot_bayes_result( phantom_profiles = results["phantom_profiles"] autocorr = results["autocorr"] - Path(figheader).mkdir(parents=True, exist_ok=True) - with open(figheader + "results.pkl", "wb") as handle: - pickle.dump(results, handle) - plt.figure() plt.plot( np.arange(0, autocorr.__len__())[np.isfinite(autocorr)], @@ -367,6 +363,7 @@ def plot_bayes_result( color="red", linestyle="dotted", ) + key = "electron_density" plot_profile( blobs[key], @@ -374,26 +371,17 @@ def plot_bayes_result( figheader=figheader, filetype=filetype, phantom_profile=phantom_profiles, + color="blue", + sharefig=True ) - key = "impurity_density" - for elem in blobs[key].element.values: - plot_profile( - blobs[key].sel(element=elem), - key, - figheader=figheader, - filename=f"{elem} density", - filetype=filetype, - phantom_profile=phantom_profiles, - color="red", - ) key = "ion_density" plot_profile( blobs[key].sel(element="h"), key, figheader=figheader, - filename=f"h density", filetype=filetype, phantom_profile=phantom_profiles, + sharefig=True, color="red", ) key = "fast_density" @@ -401,10 +389,24 @@ def plot_bayes_result( blobs[key], key, figheader=figheader, + filename="density", filetype=filetype, phantom_profile=phantom_profiles, - color="red", + color="green", ) + + key = "impurity_density" + for elem in blobs[key].element.values: + plot_profile( + blobs[key].sel(element=elem), + key, + figheader=figheader, + filename=f"{elem} density", + filetype=filetype, + phantom_profile=phantom_profiles, + color="red", + ) + key = "neutral_density" plot_profile( blobs[key], @@ -425,7 +427,7 @@ def plot_bayes_result( if __name__ == "__main__": - filehead = "./results/10009_60ms_short/" + filehead = "./results/10009_withoutCM/" with open(filehead + "results.pkl", "rb") as handle: results = pickle.load(handle) plot_bayes_result(results, filehead, filetype=".png") From 533a44a526ff533a1fd03f9d4c0a79568e39cb0e Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 22 Jun 2023 16:05:31 +0100 Subject: [PATCH 053/160] multiplot formatting --- indica/workflows/bayes_plots.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/indica/workflows/bayes_plots.py b/indica/workflows/bayes_plots.py index 3a76d9a3..403a5cf2 100644 --- a/indica/workflows/bayes_plots.py +++ b/indica/workflows/bayes_plots.py @@ -170,7 +170,7 @@ def violinplot( figheader="./results/test/", **kwargs, ): - set_plot_rcparams() + set_plot_rcparams("multi") fig, axs = plt.subplots( 1, 1, @@ -188,7 +188,8 @@ def violinplot( violin["bodies"][0].set_edgecolor("black") axs.set_xlabel(key) top = axs.get_ylim()[1] - axs.set_ylim(bottom=0, top=top * 1.1) + bot = axs.get_ylim()[0] + axs.set_ylim( top=top * 1.1, bottom=bot * 0.9) axs.set_ylabel(f"{ylabel}") y = diag_data[key].sel(t=data[key].t).values axs.errorbar( @@ -316,6 +317,8 @@ def plot_bayes_result( figheader=figheader, ylabel="Temperature [eV]", ) + + set_plot_rcparams("multi") key = "xrcs.spectra" if key in blobs.keys(): _plot_1d( @@ -427,7 +430,7 @@ def plot_bayes_result( if __name__ == "__main__": - filehead = "./results/10009_withoutCM/" + filehead = "./results/10009_60ms_long/" with open(filehead + "results.pkl", "rb") as handle: results = pickle.load(handle) plot_bayes_result(results, filehead, filetype=".png") From 751ca8139f77e0946be172dcecb09600cdef046e Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 27 Jun 2023 10:11:16 +0100 Subject: [PATCH 054/160] generalised writing kinetic profiles to blobs --- indica/bayesmodels.py | 38 ++++++++++++-------------------------- 1 file changed, 12 insertions(+), 26 deletions(-) diff --git a/indica/bayesmodels.py b/indica/bayesmodels.py index 4ef46d63..d6802160 100644 --- a/indica/bayesmodels.py +++ b/indica/bayesmodels.py @@ -6,10 +6,13 @@ from scipy.stats import uniform np.seterr(all="ignore") - warnings.simplefilter("ignore", category=FutureWarning) +PROFILES = ["electron_temperature", "electron_density", "ion_temperature", + "impurity_density", "fast_density", "neutral_density"] + + def gaussian(x, mean, sigma): return 1 / (sigma * np.sqrt(2 * np.pi)) * np.exp(-1 / 2 * ((x - mean) / sigma) ** 2) @@ -182,7 +185,7 @@ def ln_posterior(self, parameters: dict, **kwargs): """ ln_prior = self._ln_prior(parameters) - if ln_prior == -np.inf: # Don't call model if outside priors + if ln_prior == -np.inf: # Don't call models if outside priors return -np.inf, {} self.plasma.update_profiles(parameters) @@ -190,29 +193,12 @@ def ln_posterior(self, parameters: dict, **kwargs): ln_likelihood = self._ln_likelihood() # compare results to data ln_posterior = ln_likelihood + ln_prior - kin_profs = { - "electron_density": self.plasma.electron_density.sel( - t=self.plasma.time_to_calculate - ), - "electron_temperature": self.plasma.electron_temperature.sel( - t=self.plasma.time_to_calculate - ), - "ion_temperature": self.plasma.ion_temperature.sel( - t=self.plasma.time_to_calculate - ), - "impurity_density": self.plasma.impurity_density.sel( - t=self.plasma.time_to_calculate - ), - "ion_density": self.plasma.ion_density.sel( - t=self.plasma.time_to_calculate - ), - "fast_density": self.plasma.fast_density.sel( - t=self.plasma.time_to_calculate - ), - "neutral_density": self.plasma.neutral_density.sel( - t=self.plasma.time_to_calculate - ), - - } + kin_profs = {} + for profile_key in PROFILES: + if hasattr(self.plasma, profile_key): + kin_profs[profile_key] = getattr(self.plasma, profile_key).sel(t=self.plasma.time_to_calculate) + else: + raise ValueError(f"plasma does not have attribute {profile_key}") + blob = deepcopy({**self.bckc, **kin_profs}) return ln_posterior, blob From 59c8316053de2d232e41c63174b6939b6429c355 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 27 Jun 2023 11:25:43 +0100 Subject: [PATCH 055/160] function autocorr plot added --- indica/workflows/bayes_plots.py | 31 +++++++++++++++++++------------ 1 file changed, 19 insertions(+), 12 deletions(-) diff --git a/indica/workflows/bayes_plots.py b/indica/workflows/bayes_plots.py index 403a5cf2..ea943eae 100644 --- a/indica/workflows/bayes_plots.py +++ b/indica/workflows/bayes_plots.py @@ -231,6 +231,23 @@ def histograms(data, diag_data, filename): plt.close() + +def plot_autocorr(autocorr, figheader, filetype=".png"): + + plt.figure() + plt.plot( + np.arange(0, autocorr.__len__())[np.isfinite(autocorr)], + autocorr[np.isfinite(autocorr)], + label="average auto-correlation time", + ) + + plt.legend() + plt.xlabel("iterations") + plt.ylabel("auto-correlation time (iterations)") + plt.savefig(figheader + "average_tau" + filetype) + plt.close() + + def plot_bayes_result( results, figheader="./results/test/", @@ -246,17 +263,7 @@ def plot_bayes_result( phantom_profiles = results["phantom_profiles"] autocorr = results["autocorr"] - plt.figure() - plt.plot( - np.arange(0, autocorr.__len__())[np.isfinite(autocorr)], - autocorr[np.isfinite(autocorr)], - label="average tau", - ) - plt.legend() - plt.xlabel("iterations") - plt.ylabel("auto-correlation time (iterations)") - plt.savefig(figheader + "average_tau" + filetype) - plt.close() + plot_autocorr(autocorr, figheader, filetype=filetype) if "cxff_pi.ti" in blobs.keys(): blobs["cxff_pi.ti0"] = blobs["cxff_pi.ti"].sel( @@ -430,7 +437,7 @@ def plot_bayes_result( if __name__ == "__main__": - filehead = "./results/10009_60ms_long/" + filehead = "./results/test/" with open(filehead + "results.pkl", "rb") as handle: results = pickle.load(handle) plot_bayes_result(results, filehead, filetype=".png") From b2dee2d22d76d146f4f2c5836e7659f74e766f42 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 27 Jun 2023 11:26:12 +0100 Subject: [PATCH 056/160] ion density added to kinetic profiles saved --- indica/bayesmodels.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/indica/bayesmodels.py b/indica/bayesmodels.py index d6802160..fc791837 100644 --- a/indica/bayesmodels.py +++ b/indica/bayesmodels.py @@ -9,7 +9,7 @@ warnings.simplefilter("ignore", category=FutureWarning) -PROFILES = ["electron_temperature", "electron_density", "ion_temperature", +PROFILES = ["electron_temperature", "electron_density", "ion_temperature", "ion_density", "impurity_density", "fast_density", "neutral_density"] From 8eee24ebfd9516e55f2bf914764e69c796c7fd8d Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 27 Jun 2023 16:04:21 +0100 Subject: [PATCH 057/160] inital commit of abstract workflow and example of usage --- indica/workflows/abstract_bayes_workflow.py | 47 +++ indica/workflows/test_bayes_workflow.py | 336 ++++++++++++++++++++ 2 files changed, 383 insertions(+) create mode 100644 indica/workflows/abstract_bayes_workflow.py create mode 100644 indica/workflows/test_bayes_workflow.py diff --git a/indica/workflows/abstract_bayes_workflow.py b/indica/workflows/abstract_bayes_workflow.py new file mode 100644 index 00000000..95d8d6d0 --- /dev/null +++ b/indica/workflows/abstract_bayes_workflow.py @@ -0,0 +1,47 @@ +from abc import ABC, abstractmethod + + +class BayesWorkflow(ABC): + @abstractmethod + def setup_plasma(self): + """ + Contains all methods and settings for plasma object to be used in optimisation + """ + self.plasma = None + + @abstractmethod + def read_data(self, diagnostics: list): + """ + Reads data from server + + Returns + + dictionary of data + """ + self.data = {} + + def setup_opt_data(self, phantom: bool = False): + """ + Prepare the data in necessary format for optimiser + """ + self.opt_data = {} + + @abstractmethod + def setup_models(self, diagnostics: list): + """ + Initialise models normally requires data to be read so transforms can be set + + """ + self.models = {} + + @abstractmethod + def setup_optimiser(self): + """ + Initialise and provide settings for optimiser + """ + self.bayesopt = None + + @abstractmethod + def __call__(self, *args, **kwargs): + return {} + diff --git a/indica/workflows/test_bayes_workflow.py b/indica/workflows/test_bayes_workflow.py new file mode 100644 index 00000000..f7d4206d --- /dev/null +++ b/indica/workflows/test_bayes_workflow.py @@ -0,0 +1,336 @@ +import emcee +import numpy as np +import flatdict +from scipy.stats import loguniform + +from indica.readers.read_st40 import ReadST40 +from indica.bayesmodels import BayesModels, get_uniform +from indica.models.interferometry import Interferometry +from indica.models.helike_spectroscopy import Helike_spectroscopy +from indica.models.charge_exchange import ChargeExchange +from indica.models.equilibrium_reconstruction import EquilibriumReconstruction +from indica.models.plasma import Plasma +from indica.converters.line_of_sight import LineOfSightTransform + +from abstract_bayes_workflow import BayesWorkflow + +# global configurations +DEFAULT_PHANTOM_PARAMS = { + "Ne_prof.y0": 5e19, + "Ne_prof.wcenter": 0.4, + "Ne_prof.peaking": 2, + "Ne_prof.y1": 2e18, + "Ne_prof.yend": 1e18, + "Ne_prof.wped": 2, + "Nimp_prof.y0": 3e16, + "Nimp_prof.y1": 0.5e16, + "Nimp_prof.wcenter": 0.4, + "Nimp_prof.wped": 6, + "Nimp_prof.peaking": 2, + "Te_prof.y0": 3000, + "Te_prof.y1": 50, + "Te_prof.wcenter": 0.4, + "Te_prof.wped": 3, + "Te_prof.peaking": 2, + "Ti_prof.y0": 6000, + "Ti_prof.y1": 50, + "Ti_prof.wcenter": 0.4, + "Ti_prof.wped": 3, + "Ti_prof.peaking": 2, +} + +DEFAULT_PRIORS = { + "Ne_prof.y0": get_uniform(2e19, 4e20), + "Ne_prof.y1": get_uniform(1e18, 1e19), + "Ne_prof.y0/Ne_prof.y1": lambda x1, x2: np.where((x1 > x2 * 2), 1, 0), + "Ne_prof.wped": get_uniform(1, 6), + "Ne_prof.wcenter": get_uniform(0.1, 0.8), + "Ne_prof.peaking": get_uniform(1, 6), + "ar_conc": loguniform(0.0001, 0.01), + "Nimp_prof.y0": get_uniform(1e16, 1e18), + "Nimp_prof.y1": get_uniform(1e15, 2e16), + "Ne_prof.y0/Nimp_prof.y0": lambda x1, x2: np.where( + (x1 > x2 * 100) & (x1 < x2 * 1e5), 1, 0 + ), + "Nimp_prof.y0/Nimp_prof.y1": lambda x1, x2: np.where((x1 > x2), 1, 0), + "Nimp_prof.wped": get_uniform(1, 6), + "Nimp_prof.wcenter": get_uniform(0.1, 0.8), + "Nimp_prof.peaking": get_uniform(1, 20), + "Nimp_prof.peaking/Ne_prof.peaking": lambda x1, x2: np.where( + (x1 > x2), 1, 0 + ), # impurity always more peaked + "Te_prof.y0": get_uniform(1000, 5000), + "Te_prof.wped": get_uniform(1, 6), + "Te_prof.wcenter": get_uniform(0.1, 0.6), + "Te_prof.peaking": get_uniform(1, 6), + "Ti_prof.y0/Te_prof.y0": lambda x1, x2: np.where(x1 > x2, 1, 0), # hot ion mode + "Ti_prof.y0": get_uniform(3000, 10000), + "Ti_prof.wped": get_uniform(1, 6), + "Ti_prof.wcenter": get_uniform(0.1, 0.6), + "Ti_prof.peaking": get_uniform(1, 20), +} + +OPTIMISED_PARAMS = [ + # "Ne_prof.y1", + "Ne_prof.y0", + "Ne_prof.peaking", + "Ne_prof.wcenter", + "Ne_prof.wped", + # "ar_conc", + # "Nimp_prof.y1", + "Nimp_prof.y0", + # "Nimp_prof.wcenter", + # "Nimp_prof.wped", + # "Nimp_prof.peaking", + "Te_prof.y0", + # "Te_prof.wped", + # "Te_prof.wcenter", + # "Te_prof.peaking", + "Ti_prof.y0", + # "Ti_prof.wped", + # "Ti_prof.wcenter", + # "Ti_prof.peaking", +] + +OPTIMISED_QUANTITY = ["cxff_pi.ti", "efit.wp", "smmh1.ne"] + + +class TestWorkflow(BayesWorkflow): + def __init__( + self, + pulse=None, + nwalkers=50, + tstart=0.02, + tend=0.10, + dt=0.01, + tsample=0.06, + iterations=100, + burn_in_fraction=0, + + phantom=False, + diagnostics=None, + ): + self.pulse = pulse + self.tstart = tstart + self.tend = tend + self.dt = dt + self.tsample = tsample + self.nwalkers = nwalkers + self.iterations = iterations + self.burn_in_fraction = burn_in_fraction + + self.phantom = phantom + self.diagnostics = diagnostics + + self.setup_plasma() + self.read_data(diagnostics) + self.setup_opt_data(phantom=self.phantom) + self.setup_models(diagnostics) + self.setup_optimiser() + + def setup_plasma(self): + self.plasma = Plasma( + tstart=self.tstart, + tend=self.tend, + dt=self.dt, + main_ion="h", + impurities=("ar", "c"), + impurity_concentration=( + 0.001, + 0.02, + ), + full_run=False, + n_rad=20, + ) + self.plasma.Ti_ref = True + + self.plasma.time_to_calculate = self.plasma.t[ + np.abs(self.tsample - self.plasma.t).argmin() + ] + self.plasma.update_profiles(DEFAULT_PHANTOM_PARAMS) + self.plasma.build_atomic_data(calc_power_loss=False) + + def read_data(self, diagnostics: list): + self.reader = ReadST40( + self.pulse, tstart=self.tstart, tend=self.tend, dt=self.dt + ) + self.reader(diagnostics) + self.plasma.set_equilibrium(self.reader.equilibrium) + self.data = self.reader.binned_data + + def setup_opt_data(self, phantom=False): + if phantom: + self._phantom_data() + else: + self._exp_data() + + def setup_models(self, diagnostics: list): + self.models = {} + for diag in self.diagnostics: + if diag == "smmh1": + # los_transform = self.data["smmh1"]["ne"].transform + machine_dims = self.plasma.machine_dimensions + origin = np.array([[-0.38063365, 0.91893092, 0.01]]) + # end = np.array([[0, 0, 0.01]]) + direction = np.array([[0.38173721, -0.92387953, -0.02689453]]) + los_transform = LineOfSightTransform( + origin[:, 0], + origin[:, 1], + origin[:, 2], + direction[:, 0], + direction[:, 1], + direction[:, 2], + name="", + machine_dimensions=machine_dims, + passes=2, + ) + los_transform.set_equilibrium(self.plasma.equilibrium) + model = Interferometry(name=diag) + model.set_los_transform(los_transform) + model.plasma = self.plasma + + elif diag == "xrcs": + los_transform = self.data["xrcs"]["te_kw"].transform + model = Helike_spectroscopy( + name="xrcs", + window_masks=[slice(0.394, 0.396)], + window_vector=self.data[diag][ + "spectra" + ].wavelength.values + * 0.1, + ) + model.set_los_transform(los_transform) + model.plasma = self.plasma + + elif diag == "efit": + model = EquilibriumReconstruction(name="efit") + model.plasma = self.plasma + + elif diag == "cxff_pi": + transform = self.data[diag]["ti"].transform + transform.set_equilibrium(self.plasma.equilibrium) + model = ChargeExchange(name=diag, element="ar") + model.set_transect_transform(transform) + model.plasma = self.plasma + else: + raise ValueError(f"{diag} not found in setup_models") + + self.models[diag] = model + + + def setup_optimiser(self, ): + + self.bayesopt = BayesModels( + plasma=self.plasma, + data=self.opt_data, + diagnostic_models=[*self.models.values()], + quant_to_optimise=OPTIMISED_QUANTITY, + priors=DEFAULT_PRIORS, + ) + + ndim = len(OPTIMISED_PARAMS) + + self.move = [(emcee.moves.StretchMove(), 0.9), (emcee.moves.DEMove(), 0.1)] + self.sampler = emcee.EnsembleSampler( + self.nwalkers, + ndim, + log_prob_fn=self.bayesopt.ln_posterior, + parameter_names=OPTIMISED_PARAMS, + moves=self.move, + ) + + self.start_points = self.bayesopt.sample_from_priors( + OPTIMISED_PARAMS, size=self.nwalkers + ) + + + def _phantom_data(self, noise=False, noise_factor=0.1): + self.opt_data = {} + if "smmh1" in self.diagnostics: + self.opt_data["smmh1.ne"] = ( + self.models["smmh1"]() + .pop("ne") + .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) + ) + if "xrcs" in self.diagnostics: + self.opt_data["xrcs.spectra"] = ( + self.models["xrcs"]() + .pop("spectra") + .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) + ) + self.opt_data["xrcs.background"] = None + if "cxff_pi" in self.diagnostics: + cxrs_data = ( + self.models["cxff_pi"]() + .pop("ti") + .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) + ) + self.opt_data["cxff_pi.ti"] = cxrs_data.where(cxrs_data.channel == 2) + if "efit" in self.diagnostics: + self.opt_data["efit.wp"] = ( + self.models["efit"]() + .pop("wp") + .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) + ) + + if noise: + self.opt_data["smmh1.ne"] = self.opt_data["smmh1.ne"] + self.opt_data[ + "smmh1.ne" + ].max().values * np.random.normal(0, noise_factor, None) + self.opt_data["xrcs.spectra"] = self.opt_data[ + "xrcs.spectra" + ] + np.random.normal( + 0, + np.sqrt(self.opt_data["xrcs.spectra"].values[0,]), + self.opt_data["xrcs.spectra"].shape[1], + ) + self.opt_data["cxff_pi.ti"] = self.opt_data[ + "cxff_pi.ti" + ] + self.opt_data["cxff_pi.ti"].max().values * np.random.normal( + 0, noise_factor, self.opt_data["cxff_pi.ti"].shape[1] + ) + self.opt_data["efit.wp"] = self.opt_data["efit.wp"] + self.opt_data[ + "efit.wp" + ].max().values * np.random.normal(0, noise_factor, None) + + def _exp_data(self): + self.opt_data = flatdict.FlatDict(self.data, ".") + if "xrcs" in self.diagnostics: + self.opt_data["xrcs.spectra"]["wavelength"] = ( + self.opt_data["xrcs.spectra"].wavelength * 0.1 + ) + background = self.opt_data["xrcs.spectra"].where( + (self.opt_data["xrcs.spectra"].wavelength < 0.392) + & (self.opt_data["xrcs.spectra"].wavelength > 0.388), + drop=True, + ) + self.opt_data["xrcs.background"] = background.mean(dim="wavelength") + self.opt_data["xrcs.spectra"]["error"] = np.sqrt( + self.opt_data["xrcs.spectra"] + background.std(dim="wavelength") ** 2 + ) + + if "cxff_pi" in self.diagnostics: + self.opt_data["cxff_pi"]["ti"] = self.opt_data["cxff_pi"]["ti"].where( + self.opt_data["cxff_pi"]["ti"].channel == 2, drop=True + ) + + def __call__(self, *args, **kwargs): + self.sampler.run_mcmc(self.start_points, self.iterations, progress=True) + blobs = self.sampler.get_blobs( + discard=int(self.iterations * self.burn_in_fraction), flat=True + ) + return blobs + + +if __name__ == "__main__": + run = TestWorkflow( + pulse=10009, + dt=0.005, + tsample=0.060, + diagnostics=["efit", "smmh1", "cxff_pi"], + + iterations=10, + nwalkers=20, + burn_in_fraction=0.10, + ) + run() From 408b619627d4e65a7897e3504759a1ad152f559e Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 27 Jun 2023 16:04:53 +0100 Subject: [PATCH 058/160] doc strings --- indica/workflows/abstract_bayes_workflow.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/indica/workflows/abstract_bayes_workflow.py b/indica/workflows/abstract_bayes_workflow.py index 95d8d6d0..caa14bc4 100644 --- a/indica/workflows/abstract_bayes_workflow.py +++ b/indica/workflows/abstract_bayes_workflow.py @@ -16,20 +16,20 @@ def read_data(self, diagnostics: list): Returns - dictionary of data + nested dictionary of data """ self.data = {} def setup_opt_data(self, phantom: bool = False): """ - Prepare the data in necessary format for optimiser + Prepare the data in necessary format for optimiser i.e. flat dictionary """ self.opt_data = {} @abstractmethod def setup_models(self, diagnostics: list): """ - Initialise models normally requires data to be read so transforms can be set + Initialising models normally requires data to be read so transforms can be set """ self.models = {} From de237dd3c68976e740c834e79978aa7457621c16 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 27 Jun 2023 16:06:27 +0100 Subject: [PATCH 059/160] renaming workflow --- .../{test_bayes_workflow.py => bayes_workflow_example.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename indica/workflows/{test_bayes_workflow.py => bayes_workflow_example.py} (100%) diff --git a/indica/workflows/test_bayes_workflow.py b/indica/workflows/bayes_workflow_example.py similarity index 100% rename from indica/workflows/test_bayes_workflow.py rename to indica/workflows/bayes_workflow_example.py From 06c07462aea64efdfbb537ccc8629cd772ace81f Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Fri, 30 Jun 2023 13:38:43 +0100 Subject: [PATCH 060/160] Ti_ref as default --- indica/models/plasma.py | 11 ++----- ..._dev_workflow.py => bayes_workflow_dev.py} | 33 ++++++++++--------- 2 files changed, 19 insertions(+), 25 deletions(-) rename indica/workflows/{bayes_dev_workflow.py => bayes_workflow_dev.py} (96%) diff --git a/indica/models/plasma.py b/indica/models/plasma.py index 1f57439c..fcf79388 100644 --- a/indica/models/plasma.py +++ b/indica/models/plasma.py @@ -462,16 +462,9 @@ def assign_profiles(self, impurity_to_profile=None): self.electron_density.loc[dict(t=t)] = self.Ne_prof() self.calc_impurity_density(self.impurities) self.electron_temperature.loc[dict(t=t)] = self.Te_prof() - if self.Ti_ref: - self.ion_temperature.loc[dict(t=t)] = self.Ti_prof(y0_ref=self.Te_prof.y0) - else: - self.ion_temperature.loc[dict(t=t)] = self.Ti_prof() + self.ion_temperature.loc[dict(t=t)] = self.Ti_prof(y0_ref=self.Te_prof.y0) # Temp default Ti_ref self.toroidal_rotation.loc[dict(t=t)] = self.Vrot_prof() - self.neutral_density.loc[ - dict( - t=t, - ) - ] = self.Nh_prof() + self.neutral_density.loc[dict(t=t)] = self.Nh_prof() if ( impurity_to_profile is not None ): # overwrite impurity profile with non-Ne dependent profile diff --git a/indica/workflows/bayes_dev_workflow.py b/indica/workflows/bayes_workflow_dev.py similarity index 96% rename from indica/workflows/bayes_dev_workflow.py rename to indica/workflows/bayes_workflow_dev.py index d6e72780..5961bb5f 100644 --- a/indica/workflows/bayes_dev_workflow.py +++ b/indica/workflows/bayes_workflow_dev.py @@ -97,7 +97,7 @@ # "Ti_prof.peaking", ] -OPTIMISED_QUANTITY = ["xrcs.spectra", "cxff_pi.ti", "efit.wp", "smmh1.ne"] +OPTIMISED_QUANTITY = ["cxff_pi.ti", "efit.wp", "smmh1.ne"] class BayesWorkflow: @@ -161,9 +161,10 @@ def __init__( ] self.plasma.update_profiles(DEFAULT_PHANTOM_PARAMS) self.plasma.build_atomic_data(calc_power_loss=False) - self.init_fast_particles() + self.read_st40(diagnostics) + self.init_models() if self.phantom: @@ -188,12 +189,12 @@ def __init__( log_prob_fn=self.bayes_run.ln_posterior, parameter_names=OPTIMISED_PARAMS, moves=self.move, - kwargs={ - "moment_analysis": False, - "calc_spectra": True, - "minimum_lines": False, - "background": self.flat_data["xrcs.background"], - }, + # kwargs={ + # "moment_analysis": False, + # "calc_spectra": True, + # "minimum_lines": False, + # "background": self.flat_data["xrcs.background"], + # }, ) if not self.center_mass_sampling: @@ -407,7 +408,7 @@ def exp_data(self): def __call__(self, *args, **kwargs): autocorr = sample_with_autocorr( - self.sampler, self.start_points, iterations=self.iterations, auto_sample=5 + self.sampler, self.start_points, self.iterations, auto_sample=10 ) blobs = self.sampler.get_blobs( discard=int(self.iterations * self.burn_in_fraction), flat=True @@ -473,7 +474,7 @@ def __call__(self, *args, **kwargs): plot_bayes_result(self.results, figheader=self.result_path, filetype=".png") -def sample_with_autocorr(sampler, start_points, iterations=10, auto_sample=5): +def sample_with_autocorr(sampler, start_points, iterations, auto_sample=5): autocorr = np.ones((iterations,)) * np.nan old_tau = np.inf for sample in sampler.sample( @@ -483,14 +484,14 @@ def sample_with_autocorr(sampler, start_points, iterations=10, auto_sample=5): ): if sampler.iteration % auto_sample: continue - new_tau = sampler.get_autocorr_time(tol=0) - autocorr[sampler.iteration - 1] = np.mean(new_tau) + new_tau = np.mean(sampler.get_autocorr_time(tol=0)) + autocorr[sampler.iteration - 1] = new_tau converged = np.all(new_tau * 50 < sampler.iteration) converged &= np.all(np.abs(old_tau - new_tau) / new_tau < 0.01) if converged: break old_tau = new_tau - autocorr = autocorr[: sampler.iteration] + autocorr = autocorr[: sampler.iteration, ] return autocorr @@ -507,13 +508,13 @@ def sample_with_autocorr(sampler, start_points, iterations=10, auto_sample=5): pulse=10009, result_path="./results/test/", iterations=20, - nwalkers=50, + nwalkers=20, burn_in_fraction=0.10, dt=0.005, tsample=0.060, - diagnostics=["xrcs", "efit", "smmh1", "cxff_pi"], + diagnostics=["efit", "smmh1", "cxff_pi"], phantom=False, - center_mass_sampling=True, + center_mass_sampling=False, Ti_ref=True, mds_write=mds_write, save_plots=True From 5435769855061ccc964a735793840d18f1fe12a1 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Fri, 30 Jun 2023 13:39:30 +0100 Subject: [PATCH 061/160] abstract methods added --- indica/workflows/abstract_bayes_workflow.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/indica/workflows/abstract_bayes_workflow.py b/indica/workflows/abstract_bayes_workflow.py index caa14bc4..d2bb2f69 100644 --- a/indica/workflows/abstract_bayes_workflow.py +++ b/indica/workflows/abstract_bayes_workflow.py @@ -20,6 +20,8 @@ def read_data(self, diagnostics: list): """ self.data = {} + + @abstractmethod def setup_opt_data(self, phantom: bool = False): """ Prepare the data in necessary format for optimiser i.e. flat dictionary From ee3a62e1af7ef4a8cafcf4897e6d2907d89c4ccd Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 4 Jul 2023 14:08:00 +0100 Subject: [PATCH 062/160] Changed to new MDSPlus structure --- indica/workflows/bayes_plots.py | 186 +++++++++++++++++--------------- 1 file changed, 97 insertions(+), 89 deletions(-) diff --git a/indica/workflows/bayes_plots.py b/indica/workflows/bayes_plots.py index ea943eae..9d781436 100644 --- a/indica/workflows/bayes_plots.py +++ b/indica/workflows/bayes_plots.py @@ -1,9 +1,10 @@ -from pathlib import Path +import os import corner import matplotlib.pyplot as plt import numpy as np import pickle +import flatdict from indica.utilities import set_plot_rcparams, set_axis_sci @@ -22,35 +23,21 @@ def plot_profile( ): set_plot_rcparams("profiles") - if blobkey == "electron_temperature": - legkey = "Te" - elif blobkey == "ion_temperature": - legkey = "Ti" - elif blobkey == "ion_density": - legkey = "Ni" - elif blobkey == "electron_density": - legkey = "Ne" - elif blobkey == "impurity_density": - legkey = "Nimp" - elif blobkey == "fast_density": - legkey = "Nfast" - elif blobkey == "neutral_density": - legkey = "Nneut" plt.fill_between( profile.rho_poloidal, profile.quantile(0.16, dim="index"), profile.quantile(0.84, dim="index"), - label=f"{legkey}, 68% Confidence", + label=f"{blobkey}, 68% Confidence", zorder=3, color=color, alpha=0.9, ) - if legkey != "Nfast": + if blobkey != "NFAST": plt.fill_between( profile.rho_poloidal, profile.quantile(0.025, dim="index"), profile.quantile(0.975, dim="index"), - label=f"{legkey}, 95% Confidence", + label=f"{blobkey}, 95% Confidence", zorder=2, color="grey", alpha=0.4, @@ -59,19 +46,19 @@ def plot_profile( profile.rho_poloidal, profile.quantile(0.00, dim="index"), profile.quantile(1.00, dim="index"), - label=f"{legkey}, Max-Min", + label=f"{blobkey}, Max-Min", zorder=1, color="lightgrey", alpha=0.2, ) - if phantom_profile is not None: + if phantom_profile["FLAG"]: if "element" in phantom_profile[blobkey].dims: phantom = phantom_profile[blobkey].sel(element="ar") else: phantom = phantom_profile[blobkey] phantom.plot( - label=f"{legkey}, phantom profile", + label=f"{blobkey}, phantom profile", linestyle=linestyle, color="black", zorder=4, @@ -232,16 +219,17 @@ def histograms(data, diag_data, filename): -def plot_autocorr(autocorr, figheader, filetype=".png"): +def plot_autocorr(autocorr, param_names, figheader, filetype=".png"): plt.figure() + + x_data = np.ones(shape=(autocorr.shape)) * np.arange(0, autocorr[:,0].__len__())[:,None] plt.plot( - np.arange(0, autocorr.__len__())[np.isfinite(autocorr)], - autocorr[np.isfinite(autocorr)], - label="average auto-correlation time", + np.where(np.isfinite(autocorr), x_data, np.nan), + autocorr, + "x" ) - - plt.legend() + plt.legend(param_names) plt.xlabel("iterations") plt.ylabel("auto-correlation time (iterations)") plt.savefig(figheader + "average_tau" + filetype) @@ -255,28 +243,38 @@ def plot_bayes_result( **kwargs, ): - diag_data = results["diag_data"] - blobs = results["blobs"] - samples = results["samples"] - prior_samples = results["prior_samples"] - param_names = results["param_names"] - phantom_profiles = results["phantom_profiles"] - autocorr = results["autocorr"] - - plot_autocorr(autocorr, figheader, filetype=filetype) - - if "cxff_pi.ti" in blobs.keys(): - blobs["cxff_pi.ti0"] = blobs["cxff_pi.ti"].sel( - channel=diag_data["cxff_pi.ti"].channel + # delete all but pickle in directory + for root, dirs, files in os.walk(figheader): + for f in files: + if f.endswith(".pkl"): + continue + else: + print(f"Deleting {os.path.join(root, f)}") + os.remove(os.path.join(root, f)) + + diag_data = flatdict.FlatDict(results["DIAG_DATA"], ".") + model_data = flatdict.FlatDict(results["MODEL_DATA"], ".") + profiles = flatdict.FlatDict(results["PROFILE_STAT"], ".") + post_sample = results["OPTIMISATION"]["POST_SAMPLE"] + prior_sample = results["OPTIMISATION"]["PRIOR_SAMPLE"] + autocorr = results["OPTIMISATION"]["AUTOCORR"] + param_names = results["INPUT"]["PARAM_NAMES"] + phantom_profiles = flatdict.FlatDict(results["PHANTOMS"], ".") + + plot_autocorr(autocorr, param_names, figheader, filetype=filetype) + + if "CXFF_PI.TI" in model_data.keys(): + model_data["CXFF_PI.TI0"] = model_data["CXFF_PI.TI"].sel( + channel=diag_data["CXFF_PI.TI"].channel ) - diag_data["cxff_pi.ti0"] = diag_data["cxff_pi.ti"].sel( - channel=diag_data["cxff_pi.ti"].channel + diag_data["CXFF_PI.TI0"] = diag_data["CXFF_PI.TI"].sel( + channel=diag_data["CXFF_PI.TI"].channel ) - key = "cxff_pi.ti0" - if key in blobs.keys(): + key = "CXFF_PI.TI0" + if key in model_data.keys(): violinplot( - blobs, + model_data, key, diag_data, f"{key.replace('.', '_')}" + filetype, @@ -284,10 +282,10 @@ def plot_bayes_result( ylabel="Temperature [eV]", ) - key = "efit.wp" - if key in blobs.keys(): + key = "EFIT.WP" + if key in model_data.keys(): violinplot( - blobs, + model_data, key, diag_data, f"{key.replace('.', '_')}" + filetype, @@ -295,29 +293,29 @@ def plot_bayes_result( ylabel="Energy [J]", ) key = "smmh1.ne" - if key in blobs.keys(): + if key in model_data.keys(): violinplot( - blobs, + model_data, key, diag_data, f"{key.replace('.', '_')}" + filetype, figheader=figheader, ylabel=r"Line Integrated Density [$m^{-2}$]", ) - key = "xrcs.te_kw" - if key in blobs.keys(): + key = "XRCS.TE_KW" + if key in model_data.keys(): violinplot( - blobs, + model_data, key, diag_data, f"{key.replace('.', '_')}" + filetype, figheader=figheader, ylabel="Temperature [eV]", ) - key = "xrcs.ti_w" - if key in blobs.keys(): + key = "XRCS.TI_W" + if key in model_data.keys(): violinplot( - blobs, + model_data, key, diag_data, f"{key.replace('.', '_')}" + filetype, @@ -326,10 +324,10 @@ def plot_bayes_result( ) set_plot_rcparams("multi") - key = "xrcs.spectra" - if key in blobs.keys(): + key = "XRCS.SPECTRA" + if key in model_data.keys(): _plot_1d( - blobs, + model_data, key, diag_data, f"{key.replace('.', '_')}" + filetype, @@ -339,10 +337,10 @@ def plot_bayes_result( xlim=(0.394, 0.401), figsize=(6, 4), ) - key = "cxff_pi.ti" - if key in blobs.keys(): + key = "CXFF_PI.TI" + if key in model_data.keys(): _plot_1d( - blobs, + model_data, key, diag_data, f"{key.replace('.', '_')}" + filetype, @@ -351,9 +349,9 @@ def plot_bayes_result( xlabel="Channel", ) - key = "electron_temperature" + key = "TE" plot_profile( - blobs[key], + profiles[key], key, figheader=figheader, filetype=filetype, @@ -362,9 +360,9 @@ def plot_bayes_result( color="blue", linestyle="dashdot", ) - key = "ion_temperature" + key = "TI" plot_profile( - blobs[key].sel(element="ar"), + profiles[key].sel(element="ar"), key, figheader=figheader, filename="temperature", @@ -374,9 +372,9 @@ def plot_bayes_result( linestyle="dotted", ) - key = "electron_density" + key = "NE" plot_profile( - blobs[key], + profiles[key], key, figheader=figheader, filetype=filetype, @@ -384,9 +382,9 @@ def plot_bayes_result( color="blue", sharefig=True ) - key = "ion_density" + key = "NI" plot_profile( - blobs[key].sel(element="h"), + profiles[key].sel(element="h"), key, figheader=figheader, filetype=filetype, @@ -394,44 +392,54 @@ def plot_bayes_result( sharefig=True, color="red", ) - key = "fast_density" + key = "NFAST" plot_profile( - blobs[key], + profiles[key], key, figheader=figheader, - filename="density", + filename="densities", filetype=filetype, phantom_profile=phantom_profiles, color="green", ) - key = "impurity_density" - for elem in blobs[key].element.values: - plot_profile( - blobs[key].sel(element=elem), - key, - figheader=figheader, - filename=f"{elem} density", - filetype=filetype, - phantom_profile=phantom_profiles, - color="red", - ) + key = "NIMP1" + plot_profile( + profiles[key], + key, + figheader=figheader, + filename=f"ar density", + filetype=filetype, + phantom_profile=phantom_profiles, + color="red", + ) + + key = "NIMP2" + plot_profile( + profiles[key], + key, + figheader=figheader, + filename=f"c density", + filetype=filetype, + phantom_profile=phantom_profiles, + color="red", + ) - key = "neutral_density" + key = "NNEUTR" plot_profile( - blobs[key], + profiles[key], key, - filename="", + filename="neutral density", figheader=figheader, filetype=filetype, phantom_profile=phantom_profiles, logscale=True, ) - corner.corner(samples, labels=param_names) + corner.corner(post_sample, labels=param_names) plt.savefig(figheader + "posterior" + filetype) - corner.corner(prior_samples, labels=param_names) + corner.corner(prior_sample, labels=param_names) plt.savefig(figheader + "prior" + filetype) plt.close("all") From a8ca5fab7bc4e2a8013bdab68bf35d3cc1dc886e Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 4 Jul 2023 14:08:47 +0100 Subject: [PATCH 063/160] Adding default methods --- indica/workflows/abstract_bayes_workflow.py | 134 +++++++++++++++++++- 1 file changed, 133 insertions(+), 1 deletion(-) diff --git a/indica/workflows/abstract_bayes_workflow.py b/indica/workflows/abstract_bayes_workflow.py index d2bb2f69..fdf2d5bc 100644 --- a/indica/workflows/abstract_bayes_workflow.py +++ b/indica/workflows/abstract_bayes_workflow.py @@ -1,7 +1,21 @@ from abc import ABC, abstractmethod - +import numpy as np class BayesWorkflow(ABC): + def __init__(self, + phantom, + diagnostics): + + self.phantom = phantom + self.diagnostics = diagnostics + + self.setup_plasma() + self.save_phantom_profiles() + self.read_data(self.diagnostics) + self.setup_opt_data(self.phantom) + self.setup_models(self.diagnostics) + self.setup_optimiser() + @abstractmethod def setup_plasma(self): """ @@ -43,7 +57,125 @@ def setup_optimiser(self): """ self.bayesopt = None + def save_phantom_profiles(self, kinetic_profiles=None): + if kinetic_profiles is None: + kinetic_profiles = ["electron_density", "impurity_density", "electron_temperature", + "ion_temperature", "ion_density", "fast_density", "neutral_density"] + if self.phantoms: + phantom_profiles = {profile_key: getattr(self.plasma, profile_key).sel( + t=self.plasma.time_to_calculate).copy() + for profile_key in kinetic_profiles} + else: + phantom_profiles = {profile_key: getattr(self.plasma, profile_key).sel( + t=self.plasma.time_to_calculate) * 0 + for profile_key in kinetic_profiles} + + self.phantom_profiles = phantom_profiles + + def _build_result_dict(self): + + """ + TODO: xarray to numpy for MDSPlus writing (plotting / dimensions / units have to be figured out) + + + """ + + + result = {} + result["TIME"] = self.plasma.t + result["TIME_OPT"] = self.plasma.time_to_calculate + + result["INPUT"] = { + "BURN_FRAC": self.burn_frac, + "ITER": self.iterations, + "MODEL_KWARGS": "KWARGS", + "OPT_KWARGS": "KWARGS", + "NWALKERS": self.nwalkers, + "PARAM_NAMES": self.param_names, + "PULSE": self.pulse, + } + + result["GLOBAL"] = { + "TI0": 0, + "TE0": 0, + "NE0": 0, + "NI0": 0, + "TI0_ERR": 0, + "TE0_ERR": 0, + "NE0_ERR": 0, + "NI0_ERR": 0, + } + + result["PHANTOMS"] = { + "FLAG": self.phantoms, + "NE": self.phantom_profiles["electron_density"], + "TE": self.phantom_profiles["electron_temperature"], + "TI": self.phantom_profiles["ion_temperature"], + "NI": self.phantom_profiles["ion_density"], + "NNEUTR": self.phantom_profiles["neutral_density"], + "NFAST": self.phantom_profiles["fast_density"], + "NIMP1": self.phantom_profiles["impurity_density"].sel(element="ar"), # TODO: generalise + "NIMP2": self.phantom_profiles["impurity_density"].sel(element="c") + } + + result["PROFILE_STAT"] = { + "SAMPLES": self.samples, + "RHO_POLOIDAL": self.plasma.rho, + "NE": self.blobs["electron_density"], + "NI": self.blobs["ion_density"], + "TE": self.blobs["electron_temperature"], + "TI": self.blobs["ion_temperature"], + "NFAST": self.blobs["fast_density"], + "NNEUTR": self.blobs["neutral_density"], + "NIMP1": self.blobs["impurity_density"].sel(element="ar"), # TODO: generalise + "NIMP2": self.blobs["impurity_density"].sel(element="c") + } + + result["OPTIMISATION"] = { + "OPT_QUANTITY": self.opt_quantity, + "ACCEPT_FRAC": self.accept_frac, + "PRIOR_SAMPLE": self.prior_sample, + "POST_SAMPLE": self.post_sample, + "AUTOCORR": self.autocorr, + } + + quant_list = [item.split(".") for item in self.opt_quantity] + result["MODEL_DATA"] = {diag_name.upper(): + {quantity[1].upper(): self.blobs[f"{quantity[0]}.{quantity[1]}"] + for quantity in quant_list if quantity[0] == diag_name} + for diag_name in self.diagnostics + } + result["MODEL_DATA"]["SAMPLES"] = self.samples + + result["DIAG_DATA"] = {diag_name.upper(): + {quantity[1].upper(): self.opt_data[f"{quantity[0]}.{quantity[1]}"] + for quantity in quant_list if quantity[0] == diag_name} + for diag_name in self.diagnostics + } + self.result = result + return self.result + @abstractmethod def __call__(self, *args, **kwargs): return {} + +def sample_with_autocorr(sampler, start_points, iterations, n_params, auto_sample=5): + autocorr = np.ones(shape=(iterations, n_params)) * np.nan + old_tau = np.inf + for sample in sampler.sample( + start_points, + iterations=iterations, + progress=True, + ): + if sampler.iteration % auto_sample: + continue + new_tau = sampler.get_autocorr_time(tol=0) + autocorr[sampler.iteration - 1, ] = new_tau + converged = np.all(new_tau * 50 < sampler.iteration) + converged &= np.all(np.abs(old_tau - new_tau) / new_tau < 0.01) + if converged: + break + old_tau = new_tau + autocorr = autocorr[: sampler.iteration, ] + return autocorr \ No newline at end of file From 6f1e3308ea0ddebe3dbf7c93a467c2ca9a3ac93e Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 4 Jul 2023 15:09:48 +0100 Subject: [PATCH 064/160] moved sampling to it's own function in abstract class --- indica/workflows/bayes_workflow_example.py | 84 ++++++++++++++-------- 1 file changed, 55 insertions(+), 29 deletions(-) diff --git a/indica/workflows/bayes_workflow_example.py b/indica/workflows/bayes_workflow_example.py index f7d4206d..4bc2e39a 100644 --- a/indica/workflows/bayes_workflow_example.py +++ b/indica/workflows/bayes_workflow_example.py @@ -1,7 +1,12 @@ +import pickle + import emcee import numpy as np +import xarray as xr +import pandas as pd import flatdict from scipy.stats import loguniform +from pathlib import Path from indica.readers.read_st40 import ReadST40 from indica.bayesmodels import BayesModels, get_uniform @@ -12,9 +17,11 @@ from indica.models.plasma import Plasma from indica.converters.line_of_sight import LineOfSightTransform -from abstract_bayes_workflow import BayesWorkflow +from abstract_bayes_workflow import BayesWorkflow, sample_with_autocorr # global configurations +from indica.workflows.bayes_plots import plot_bayes_result + DEFAULT_PHANTOM_PARAMS = { "Ne_prof.y0": 5e19, "Ne_prof.wcenter": 0.4, @@ -73,9 +80,9 @@ OPTIMISED_PARAMS = [ # "Ne_prof.y1", "Ne_prof.y0", - "Ne_prof.peaking", - "Ne_prof.wcenter", - "Ne_prof.wped", + # "Ne_prof.peaking", + # "Ne_prof.wcenter", + # "Ne_prof.wped", # "ar_conc", # "Nimp_prof.y1", "Nimp_prof.y0", @@ -99,32 +106,40 @@ class TestWorkflow(BayesWorkflow): def __init__( self, pulse=None, + param_names=None, + opt_quantity=None, + nwalkers=50, tstart=0.02, tend=0.10, dt=0.01, tsample=0.06, iterations=100, - burn_in_fraction=0, + burn_frac=0, - phantom=False, + phantoms=False, diagnostics=None, ): self.pulse = pulse + self.param_names = param_names + self.opt_quantity = opt_quantity + self.tstart = tstart self.tend = tend self.dt = dt self.tsample = tsample self.nwalkers = nwalkers self.iterations = iterations - self.burn_in_fraction = burn_in_fraction + self.burn_frac = burn_frac - self.phantom = phantom + + self.phantoms = phantoms self.diagnostics = diagnostics self.setup_plasma() + self.save_phantom_profiles() self.read_data(diagnostics) - self.setup_opt_data(phantom=self.phantom) + self.setup_opt_data(phantoms=self.phantoms) self.setup_models(diagnostics) self.setup_optimiser() @@ -142,14 +157,13 @@ def setup_plasma(self): full_run=False, n_rad=20, ) - self.plasma.Ti_ref = True - self.plasma.time_to_calculate = self.plasma.t[ np.abs(self.tsample - self.plasma.t).argmin() ] self.plasma.update_profiles(DEFAULT_PHANTOM_PARAMS) self.plasma.build_atomic_data(calc_power_loss=False) + def read_data(self, diagnostics: list): self.reader = ReadST40( self.pulse, tstart=self.tstart, tend=self.tend, dt=self.dt @@ -158,8 +172,8 @@ def read_data(self, diagnostics: list): self.plasma.set_equilibrium(self.reader.equilibrium) self.data = self.reader.binned_data - def setup_opt_data(self, phantom=False): - if phantom: + def setup_opt_data(self, phantoms=False): + if phantoms: self._phantom_data() else: self._exp_data() @@ -217,7 +231,6 @@ def setup_models(self, diagnostics: list): self.models[diag] = model - def setup_optimiser(self, ): self.bayesopt = BayesModels( @@ -243,7 +256,6 @@ def setup_optimiser(self, ): OPTIMISED_PARAMS, size=self.nwalkers ) - def _phantom_data(self, noise=False, noise_factor=0.1): self.opt_data = {} if "smmh1" in self.diagnostics: @@ -278,21 +290,28 @@ def _phantom_data(self, noise=False, noise_factor=0.1): "smmh1.ne" ].max().values * np.random.normal(0, noise_factor, None) self.opt_data["xrcs.spectra"] = self.opt_data[ - "xrcs.spectra" - ] + np.random.normal( + "xrcs.spectra" + ] + np.random.normal( 0, np.sqrt(self.opt_data["xrcs.spectra"].values[0,]), self.opt_data["xrcs.spectra"].shape[1], ) self.opt_data["cxff_pi.ti"] = self.opt_data[ - "cxff_pi.ti" - ] + self.opt_data["cxff_pi.ti"].max().values * np.random.normal( + "cxff_pi.ti" + ] + self.opt_data["cxff_pi.ti"].max().values * np.random.normal( 0, noise_factor, self.opt_data["cxff_pi.ti"].shape[1] ) self.opt_data["efit.wp"] = self.opt_data["efit.wp"] + self.opt_data[ "efit.wp" ].max().values * np.random.normal(0, noise_factor, None) + self.phantom_profiles = {} + for key in ["electron_density", "impurity_density", "electron_temperature", "ion_temperature", + "ion_density", "fast_density", "neutral_density"]: + self.phantom_profiles[key] = getattr(self.plasma, key).sel( + t=self.plasma.time_to_calculate + ).copy() + def _exp_data(self): self.opt_data = flatdict.FlatDict(self.data, ".") if "xrcs" in self.diagnostics: @@ -314,12 +333,12 @@ def _exp_data(self): self.opt_data["cxff_pi"]["ti"].channel == 2, drop=True ) - def __call__(self, *args, **kwargs): - self.sampler.run_mcmc(self.start_points, self.iterations, progress=True) - blobs = self.sampler.get_blobs( - discard=int(self.iterations * self.burn_in_fraction), flat=True - ) - return blobs + def __call__(self, filepath = "./results/test/", **kwargs): + self.run_sampler() + self.save_pickle(filepath=filepath) + plot_bayes_result(self.result, filepath) + return self.result + if __name__ == "__main__": @@ -329,8 +348,15 @@ def __call__(self, *args, **kwargs): tsample=0.060, diagnostics=["efit", "smmh1", "cxff_pi"], - iterations=10, - nwalkers=20, - burn_in_fraction=0.10, + iterations=20, + nwalkers=10, + burn_frac=0.10, + param_names = OPTIMISED_PARAMS, + opt_quantity=OPTIMISED_QUANTITY ) - run() + + test = run(filepath="./results/test/") + + flattest = flatdict.FlatDict(test, delimiter=".") + for key in flattest.keys(): + print(key) From dd3abdf8e2f82b6606cb716f2088497fafba44f0 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 4 Jul 2023 15:10:05 +0100 Subject: [PATCH 065/160] moved sampling to it's own function in abstract class --- indica/workflows/abstract_bayes_workflow.py | 42 ++++++++++++++++++++- 1 file changed, 40 insertions(+), 2 deletions(-) diff --git a/indica/workflows/abstract_bayes_workflow.py b/indica/workflows/abstract_bayes_workflow.py index fdf2d5bc..6f6a3af4 100644 --- a/indica/workflows/abstract_bayes_workflow.py +++ b/indica/workflows/abstract_bayes_workflow.py @@ -1,5 +1,9 @@ from abc import ABC, abstractmethod import numpy as np +import xarray as xr +import pandas as pd +from pathlib import Path +import pickle class BayesWorkflow(ABC): def __init__(self, @@ -155,9 +159,43 @@ def _build_result_dict(self): self.result = result return self.result + def run_sampler(self): + + self.autocorr = sample_with_autocorr( + self.sampler, self.start_points, self.iterations, self.param_names.__len__(), auto_sample=10 + ) + blobs = self.sampler.get_blobs( + discard=int(self.iterations * self.burn_frac), flat=True + ) + blob_names = self.sampler.get_blobs().flatten()[0].keys() + self.samples = np.arange(0, blobs.__len__()) + + self.blobs = { + blob_name: xr.concat( + [data[blob_name] for data in blobs], + dim=pd.Index(self.samples, name="index"), + ) + for blob_name in blob_names + } + self.accept_frac = self.sampler.acceptance_fraction.sum() + self.prior_sample = self.bayesopt.sample_from_priors( + self.param_names, size=int(1e4) + ) + self.post_sample = self.sampler.get_chain(flat=True) + self.result = self._build_result_dict() + + + def save_pickle(self, filepath): + if filepath: + Path(filepath).mkdir(parents=True, exist_ok=True) + with open(filepath + "results.pkl", "wb") as handle: + pickle.dump(self.result, handle) + @abstractmethod - def __call__(self, *args, **kwargs): - return {} + def __call__(self, filepath="./results/test/", **kwargs): + self.run_sampler() + self.save_pickle(filepath=filepath) + return self.result def sample_with_autocorr(sampler, start_points, iterations, n_params, auto_sample=5): From 6f9d4a83cba071a719ac9108605c7209bba05383 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 4 Jul 2023 15:18:53 +0100 Subject: [PATCH 066/160] renamed abstract class to AbstractBayesWorkflow --- indica/workflows/abstract_bayes_workflow.py | 2 +- indica/workflows/bayes_workflow_example.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/indica/workflows/abstract_bayes_workflow.py b/indica/workflows/abstract_bayes_workflow.py index 6f6a3af4..4f426d48 100644 --- a/indica/workflows/abstract_bayes_workflow.py +++ b/indica/workflows/abstract_bayes_workflow.py @@ -5,7 +5,7 @@ from pathlib import Path import pickle -class BayesWorkflow(ABC): +class AbstractBayesWorkflow(ABC): def __init__(self, phantom, diagnostics): diff --git a/indica/workflows/bayes_workflow_example.py b/indica/workflows/bayes_workflow_example.py index 4bc2e39a..569d23f1 100644 --- a/indica/workflows/bayes_workflow_example.py +++ b/indica/workflows/bayes_workflow_example.py @@ -17,7 +17,7 @@ from indica.models.plasma import Plasma from indica.converters.line_of_sight import LineOfSightTransform -from abstract_bayes_workflow import BayesWorkflow, sample_with_autocorr +from abstract_bayes_workflow import AbstractBayesWorkflow # global configurations from indica.workflows.bayes_plots import plot_bayes_result @@ -102,7 +102,7 @@ OPTIMISED_QUANTITY = ["cxff_pi.ti", "efit.wp", "smmh1.ne"] -class TestWorkflow(BayesWorkflow): +class TestWorkflow(AbstractBayesWorkflow): def __init__( self, pulse=None, From 55e2f8fe3d53465a7bb7f05fbf6abea8e09afeba Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 4 Jul 2023 15:40:46 +0100 Subject: [PATCH 067/160] comments --- indica/workflows/abstract_bayes_workflow.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/indica/workflows/abstract_bayes_workflow.py b/indica/workflows/abstract_bayes_workflow.py index 4f426d48..c0025b88 100644 --- a/indica/workflows/abstract_bayes_workflow.py +++ b/indica/workflows/abstract_bayes_workflow.py @@ -161,6 +161,16 @@ def _build_result_dict(self): def run_sampler(self): + """ + TODO: unsure if keeping in abstract class is best practice + + Runs the sampler and saves certain attributes from the sampler + + Returns + + result in MDSPlus node formatting + + """ self.autocorr = sample_with_autocorr( self.sampler, self.start_points, self.iterations, self.param_names.__len__(), auto_sample=10 ) From d3bf315e2a13d7124d7dad82b3af68dd443d99d6 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 4 Jul 2023 15:58:54 +0100 Subject: [PATCH 068/160] added high density sampling as method --- indica/bayesmodels.py | 32 ++++++++++++++++++++ indica/workflows/bayes_workflow_example.py | 35 ++++++++++++++-------- 2 files changed, 54 insertions(+), 13 deletions(-) diff --git a/indica/bayesmodels.py b/indica/bayesmodels.py index fc791837..d6092f09 100644 --- a/indica/bayesmodels.py +++ b/indica/bayesmodels.py @@ -165,6 +165,38 @@ def sample_from_priors(self, param_names, size=10): samples = samples[:, 0:size] return samples.transpose() + def sample_from_high_density_region(self, param_names: list, sampler: object, nwalkers: int, nsamples = 100): + start_points = self.sample_from_priors( + param_names, size=nsamples + ) + + ln_prob, _ = sampler.compute_log_prob(start_points) + num_best_points = int(nsamples * 0.05) + index_best_start = np.argsort(ln_prob)[-num_best_points:] + best_start_points = start_points[index_best_start, :] + best_points_std = np.std(best_start_points, axis=0) + + # Passing samples through ln_prior and redrawing if they fail + samples = np.empty((param_names.__len__(), 0)) + while samples.size < param_names.__len__() * nwalkers: + sample = np.random.normal( + np.mean(best_start_points, axis=0), + best_points_std * 2, + size=(nwalkers * 5, len(param_names)), + ) + start = { + name: sample[:, idx] for idx, name in enumerate(param_names) + } + ln_prior = self._ln_prior(start) + # Convert from dictionary of arrays -> array, + # then filtering out where ln_prior is -infinity + accepted_samples = np.array(list(start.values()))[ + :, ln_prior != -np.inf + ] + samples = np.append(samples, accepted_samples, axis=1) + start_points = samples[:, 0: nwalkers].transpose() + return start_points + def ln_posterior(self, parameters: dict, **kwargs): """ Posterior probability given to optimisers diff --git a/indica/workflows/bayes_workflow_example.py b/indica/workflows/bayes_workflow_example.py index 569d23f1..4386642a 100644 --- a/indica/workflows/bayes_workflow_example.py +++ b/indica/workflows/bayes_workflow_example.py @@ -102,12 +102,13 @@ OPTIMISED_QUANTITY = ["cxff_pi.ti", "efit.wp", "smmh1.ne"] -class TestWorkflow(AbstractBayesWorkflow): +class ExampleWorkflow(AbstractBayesWorkflow): def __init__( self, pulse=None, param_names=None, opt_quantity=None, + priors = None, nwalkers=50, tstart=0.02, @@ -119,10 +120,12 @@ def __init__( phantoms=False, diagnostics=None, + sample_high_density = False, ): self.pulse = pulse self.param_names = param_names self.opt_quantity = opt_quantity + self.priors = priors self.tstart = tstart self.tend = tend @@ -135,6 +138,7 @@ def __init__( self.phantoms = phantoms self.diagnostics = diagnostics + self.sample_high_density = sample_high_density self.setup_plasma() self.save_phantom_profiles() @@ -237,24 +241,26 @@ def setup_optimiser(self, ): plasma=self.plasma, data=self.opt_data, diagnostic_models=[*self.models.values()], - quant_to_optimise=OPTIMISED_QUANTITY, - priors=DEFAULT_PRIORS, + quant_to_optimise=self.opt_quantity, + priors=self.priors, ) - ndim = len(OPTIMISED_PARAMS) - + ndim = len(self.param_names) self.move = [(emcee.moves.StretchMove(), 0.9), (emcee.moves.DEMove(), 0.1)] self.sampler = emcee.EnsembleSampler( self.nwalkers, ndim, log_prob_fn=self.bayesopt.ln_posterior, - parameter_names=OPTIMISED_PARAMS, + parameter_names=self.param_names, moves=self.move, ) - self.start_points = self.bayesopt.sample_from_priors( - OPTIMISED_PARAMS, size=self.nwalkers - ) + if self.sample_high_density: + self.start_points = self.bayesopt.sample_from_high_density_region(self.param_names, self.sampler, self.nwalkers) + else: + self.start_points = self.bayesopt.sample_from_priors( + self.param_names, size=self.nwalkers + ) def _phantom_data(self, noise=False, noise_factor=0.1): self.opt_data = {} @@ -342,17 +348,20 @@ def __call__(self, filepath = "./results/test/", **kwargs): if __name__ == "__main__": - run = TestWorkflow( + run = ExampleWorkflow( pulse=10009, dt=0.005, tsample=0.060, diagnostics=["efit", "smmh1", "cxff_pi"], + param_names=OPTIMISED_PARAMS, + opt_quantity=OPTIMISED_QUANTITY, + priors=DEFAULT_PRIORS, + iterations=20, - nwalkers=10, + nwalkers=50, burn_frac=0.10, - param_names = OPTIMISED_PARAMS, - opt_quantity=OPTIMISED_QUANTITY + sample_high_density=True ) test = run(filepath="./results/test/") From 430a44bee56c472bc6b0ec3c35c40c38c6c2474b Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 4 Jul 2023 16:18:04 +0100 Subject: [PATCH 069/160] added attributes to __init__ --- indica/workflows/abstract_bayes_workflow.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/indica/workflows/abstract_bayes_workflow.py b/indica/workflows/abstract_bayes_workflow.py index c0025b88..fae9a283 100644 --- a/indica/workflows/abstract_bayes_workflow.py +++ b/indica/workflows/abstract_bayes_workflow.py @@ -6,17 +6,25 @@ import pickle class AbstractBayesWorkflow(ABC): + @abstractmethod def __init__(self, - phantom, - diagnostics): - - self.phantom = phantom + phantoms = None, + diagnostics = None, + param_names=None, + opt_quantity=None, + priors=None, + ): + + self.phantoms = phantoms self.diagnostics = diagnostics + self.param_names = param_names + self.opt_quantity = opt_quantity + self.priors = priors self.setup_plasma() self.save_phantom_profiles() self.read_data(self.diagnostics) - self.setup_opt_data(self.phantom) + self.setup_opt_data(self.phantoms) self.setup_models(self.diagnostics) self.setup_optimiser() From d633ae5e3fec78f2e203068e42b8fc363ebe9d7d Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 4 Jul 2023 16:30:24 +0100 Subject: [PATCH 070/160] error handling __init__ --- indica/workflows/bayes_workflow_example.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/indica/workflows/bayes_workflow_example.py b/indica/workflows/bayes_workflow_example.py index 4386642a..b7677e18 100644 --- a/indica/workflows/bayes_workflow_example.py +++ b/indica/workflows/bayes_workflow_example.py @@ -109,6 +109,7 @@ def __init__( param_names=None, opt_quantity=None, priors = None, + diagnostics=None, nwalkers=50, tstart=0.02, @@ -119,27 +120,30 @@ def __init__( burn_frac=0, phantoms=False, - diagnostics=None, sample_high_density = False, ): + self.pulse = pulse self.param_names = param_names self.opt_quantity = opt_quantity self.priors = priors + self.diagnostics = diagnostics self.tstart = tstart self.tend = tend self.dt = dt self.tsample = tsample + self.nwalkers = nwalkers self.iterations = iterations self.burn_frac = burn_frac - - self.phantoms = phantoms - self.diagnostics = diagnostics self.sample_high_density = sample_high_density + for attribute in ["pulse", "param_names", "opt_quantity", "priors", "diagnostics"]: + if getattr(self, attribute) is None: + raise ValueError(f"{attribute} needs to be defined") + self.setup_plasma() self.save_phantom_profiles() self.read_data(diagnostics) @@ -352,8 +356,8 @@ def __call__(self, filepath = "./results/test/", **kwargs): pulse=10009, dt=0.005, tsample=0.060, - diagnostics=["efit", "smmh1", "cxff_pi"], - + # diagnostics=["efit", "smmh1", "cxff_pi"], + diagnostics=None, param_names=OPTIMISED_PARAMS, opt_quantity=OPTIMISED_QUANTITY, priors=DEFAULT_PRIORS, From 35b64cf36f0d6918117cb239c5a0eac2c97db2b6 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 4 Jul 2023 16:30:37 +0100 Subject: [PATCH 071/160] error handling __init__ --- indica/workflows/bayes_workflow_dev.py | 469 +++++++++---------------- 1 file changed, 170 insertions(+), 299 deletions(-) diff --git a/indica/workflows/bayes_workflow_dev.py b/indica/workflows/bayes_workflow_dev.py index 5961bb5f..ac6af031 100644 --- a/indica/workflows/bayes_workflow_dev.py +++ b/indica/workflows/bayes_workflow_dev.py @@ -1,11 +1,7 @@ import emcee import numpy as np -import xarray as xr -import pandas as pd import flatdict from scipy.stats import loguniform -from pathlib import Path -import pickle from indica.readers.read_st40 import ReadST40 from indica.bayesmodels import BayesModels, get_uniform @@ -17,6 +13,7 @@ from indica.models.plasma import Plasma from indica.converters.line_of_sight import LineOfSightTransform +from indica.workflows.abstract_bayes_workflow import AbstractBayesWorkflow from indica.writers.bda_tree import create_nodes, write_nodes # global configurations @@ -78,9 +75,9 @@ OPTIMISED_PARAMS = [ # "Ne_prof.y1", "Ne_prof.y0", - "Ne_prof.peaking", - "Ne_prof.wcenter", - "Ne_prof.wped", + # "Ne_prof.peaking", + # "Ne_prof.wcenter", + # "Ne_prof.wped", # "ar_conc", # "Nimp_prof.y1", "Nimp_prof.y0", @@ -100,140 +97,83 @@ OPTIMISED_QUANTITY = ["cxff_pi.ti", "efit.wp", "smmh1.ne"] -class BayesWorkflow: +class DevBayesWorkflow(AbstractBayesWorkflow): def __init__( - self, - pulse=None, - pulse_to_write = 23000101, - result_path="./results/example/", - nwalkers=50, - tstart=0.02, - tend=0.10, - dt=0.01, - tsample=0.06, - iterations=100, - burn_in_fraction=0, - - diagnostics=None, - center_mass_sampling=True, - Ti_ref=True, - profiles=None, - phantom=True, - mds_write=False, - save_plots=True, + self, + pulse=None, + pulse_to_write=None, + param_names=None, + opt_quantity=None, + priors = None, + + nwalkers=50, + tstart=0.02, + tend=0.10, + dt=0.01, + tsample=0.06, + iterations=100, + burn_frac=0, + + phantoms=False, + diagnostics=None, + sample_high_density = False, + mds_write=False, + fast_particles = False, ): self.pulse = pulse self.pulse_to_write = pulse_to_write + self.param_names = param_names + self.opt_quantity = opt_quantity + self.priors = priors + self.tstart = tstart self.tend = tend self.dt = dt self.tsample = tsample - self.result_path = result_path self.nwalkers = nwalkers self.iterations = iterations - self.burn_in_fraction = burn_in_fraction - self.center_mass_sampling = center_mass_sampling - self.Ti_ref = Ti_ref + self.burn_frac = burn_frac + self.phantoms = phantoms self.diagnostics = diagnostics - self.phantom = phantom - self.profiles = profiles + self.sample_high_density = sample_high_density self.mds_write = mds_write - self.save_plots = save_plots + self.fast_particles = fast_particles + + for attribute in ["pulse", "param_names", "opt_quantity", "priors", "diagnostics"]: + if getattr(self, attribute) is None: + raise ValueError(f"{attribute} needs to be defined") + + self.setup_plasma() + self.save_phantom_profiles() + self.read_data(diagnostics) + self.setup_opt_data(phantoms=self.phantoms) + self.setup_models(diagnostics) + self.setup_optimiser() + def setup_plasma(self): self.plasma = Plasma( - tstart=tstart, - tend=tend, - dt=dt, + tstart=self.tstart, + tend=self.tend, + dt=self.dt, main_ion="h", impurities=("ar", "c"), impurity_concentration=( 0.001, - 0.04, + 0.02, ), full_run=False, n_rad=20, ) - self.plasma.Ti_ref = self.Ti_ref - self.tsample = tsample self.plasma.time_to_calculate = self.plasma.t[ - np.abs(tsample - self.plasma.t).argmin() + np.abs(self.tsample - self.plasma.t).argmin() ] self.plasma.update_profiles(DEFAULT_PHANTOM_PARAMS) self.plasma.build_atomic_data(calc_power_loss=False) - self.init_fast_particles() - - self.read_st40(diagnostics) - - self.init_models() - - if self.phantom: - self.phantom_data() - self.save_profiles() - else: - self.exp_data() - - self.bayes_run = BayesModels( - plasma=self.plasma, - data=self.flat_data, - diagnostic_models=[*self.models.values()], - quant_to_optimise=OPTIMISED_QUANTITY, - priors=DEFAULT_PRIORS, - ) - - ndim = len(OPTIMISED_PARAMS) - self.move = [(emcee.moves.StretchMove(), 0.9), (emcee.moves.DEMove(), 0.1)] - self.sampler = emcee.EnsembleSampler( - nwalkers, - ndim, - log_prob_fn=self.bayes_run.ln_posterior, - parameter_names=OPTIMISED_PARAMS, - moves=self.move, - # kwargs={ - # "moment_analysis": False, - # "calc_spectra": True, - # "minimum_lines": False, - # "background": self.flat_data["xrcs.background"], - # }, - ) + if self.fast_particles: + self._init_fast_particles() - if not self.center_mass_sampling: - self.start_points = self.bayes_run.sample_from_priors( - OPTIMISED_PARAMS, size=self.nwalkers - ) - else: - # gaussian mass around most probable starting points TODO: move this function to it's own method - nsamples = 100 - start_points = self.bayes_run.sample_from_priors( - OPTIMISED_PARAMS, size=nsamples - ) - ln_prob, _ = self.sampler.compute_log_prob(start_points) - num_best_points = int(nsamples * 0.05) - index_best_start = np.argsort(ln_prob)[-num_best_points:] - best_start_points = start_points[index_best_start, :] - best_points_std = np.std(best_start_points, axis=0) - - # Passing samples through ln_prior and redrawing if they fail - samples = np.empty((OPTIMISED_PARAMS.__len__(), 0)) - while samples.size < OPTIMISED_PARAMS.__len__() * self.nwalkers: - sample = np.random.normal( - np.mean(best_start_points, axis=0), - best_points_std * 2, - size=(self.nwalkers * 5, len(OPTIMISED_PARAMS)), - ) - start = { - name: sample[:, idx] for idx, name in enumerate(OPTIMISED_PARAMS) - } - ln_prior = self.bayes_run._ln_prior(start) - # Convert from dictionary of arrays -> array, - # then filtering out where ln_prior is -infinity - accepted_samples = np.array(list(start.values()))[ - :, ln_prior != -np.inf - ] - samples = np.append(samples, accepted_samples, axis=1) - self.start_points = samples[:, 0 : self.nwalkers].transpose() - - def init_fast_particles(self): + def _init_fast_particles(self): st40_code = ReadST40(13110009, self.tstart, self.tend, dt=self.dt, tree="astra") st40_code.get_raw_data("", "astra", "RUN573") st40_code.bin_data_in_time(["astra"], self.tstart, self.tend, self.dt) @@ -253,41 +193,25 @@ def init_fast_particles(self): Pbper = code_data["pbper"].interp(rho_poloidal=self.plasma.rho, t=self.plasma.t) self.plasma.pressure_fast_perpendicular.values = Pbper.values - def read_st40(self, diagnostics=None): - self.ST40_data = ReadST40( + def read_data(self, diagnostics: list): + self.reader = ReadST40( self.pulse, tstart=self.tstart, tend=self.tend, dt=self.dt ) - self.ST40_data(diagnostics) - self.plasma.set_equilibrium(self.ST40_data.equilibrium) + self.reader(diagnostics) + self.plasma.set_equilibrium(self.reader.equilibrium) + self.data = self.reader.binned_data - def save_profiles(self): - phantom_profiles = { - "electron_density": self.plasma.electron_density.sel( - t=self.plasma.time_to_calculate - ).copy(), - "impurity_density": self.plasma.impurity_density.sel( - t=self.plasma.time_to_calculate - ).copy(), - "electron_temperature": self.plasma.electron_temperature.sel( - t=self.plasma.time_to_calculate - ).copy(), - "ion_temperature": self.plasma.ion_temperature.sel( - t=self.plasma.time_to_calculate - ).copy(), - "fast_density": self.plasma.fast_density.sel( - t=self.plasma.time_to_calculate - ).copy(), - "neutral_density": self.plasma.neutral_density.sel( - t=self.plasma.time_to_calculate - ).copy(), - } - self.profiles = phantom_profiles + def setup_opt_data(self, phantoms=False): + if phantoms: + self._phantom_data() + else: + self._exp_data() - def init_models(self): + def setup_models(self, diagnostics: list): self.models = {} for diag in self.diagnostics: if diag == "smmh1": - # los_transform = self.ST40_data.binned_data["smmh1"]["ne"].transform + # los_transform = self.data["smmh1"]["ne"].transform machine_dims = self.plasma.machine_dimensions origin = np.array([[-0.38063365, 0.91893092, 0.01]]) # end = np.array([[0, 0, 0.01]]) @@ -304,222 +228,169 @@ def init_models(self): passes=2, ) los_transform.set_equilibrium(self.plasma.equilibrium) - model = Interferometry(name="smmh1") + model = Interferometry(name=diag) model.set_los_transform(los_transform) model.plasma = self.plasma - self.models["smmh1"] = model - if "xrcs" in self.diagnostics: - los_transform = self.ST40_data.binned_data["xrcs"]["te_kw"].transform + elif diag == "xrcs": + los_transform = self.data["xrcs"]["te_kw"].transform model = Helike_spectroscopy( name="xrcs", window_masks=[slice(0.394, 0.396)], - window_vector=self.ST40_data.binned_data["xrcs"][ - "spectra" - ].wavelength.values - * 0.1, + window_vector=self.data[diag][ + "spectra" + ].wavelength.values + * 0.1, ) model.set_los_transform(los_transform) model.plasma = self.plasma - self.models["xrcs"] = model - if "efit" in self.diagnostics: + elif diag == "efit": model = EquilibriumReconstruction(name="efit") model.plasma = self.plasma - self.models["efit"] = model - if "cxff_pi" in self.diagnostics: - transform = self.ST40_data.binned_data["cxff_pi"]["ti"].transform - transform.set_equilibrium(self.ST40_data.equilibrium) - model = ChargeExchange(name="cxff_pi", element="ar") + elif diag == "cxff_pi": + transform = self.data[diag]["ti"].transform + transform.set_equilibrium(self.plasma.equilibrium) + model = ChargeExchange(name=diag, element="ar") model.set_transect_transform(transform) model.plasma = self.plasma - self.models["cxff_pi"] = model + else: + raise ValueError(f"{diag} not found in setup_models") + + self.models[diag] = model + + def setup_optimiser(self, ): + + self.bayesopt = BayesModels( + plasma=self.plasma, + data=self.opt_data, + diagnostic_models=[*self.models.values()], + quant_to_optimise=self.opt_quantity, + priors=self.priors, + ) + + ndim = len(self.param_names) + self.move = [(emcee.moves.StretchMove(), 0.9), (emcee.moves.DEMove(), 0.1)] + self.sampler = emcee.EnsembleSampler( + self.nwalkers, + ndim, + log_prob_fn=self.bayesopt.ln_posterior, + parameter_names=self.param_names, + moves=self.move, + ) - def phantom_data(self, noise=False, noise_factor=0.1): - self.flat_data = {} + if self.sample_high_density: + self.start_points = self.bayesopt.sample_from_high_density_region(self.param_names, self.sampler, self.nwalkers) + else: + self.start_points = self.bayesopt.sample_from_priors( + self.param_names, size=self.nwalkers + ) + + def _phantom_data(self, noise=False, noise_factor=0.1): + self.opt_data = {} if "smmh1" in self.diagnostics: - self.flat_data["smmh1.ne"] = ( + self.opt_data["smmh1.ne"] = ( self.models["smmh1"]() - .pop("ne") - .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) + .pop("ne") + .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) ) if "xrcs" in self.diagnostics: - self.flat_data["xrcs.spectra"] = ( + self.opt_data["xrcs.spectra"] = ( self.models["xrcs"]() - .pop("spectra") - .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) + .pop("spectra") + .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) ) - self.flat_data["xrcs.background"] = None + self.opt_data["xrcs.background"] = None if "cxff_pi" in self.diagnostics: cxrs_data = ( self.models["cxff_pi"]() - .pop("ti") - .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) + .pop("ti") + .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) ) - self.flat_data["cxff_pi.ti"] = cxrs_data.where(cxrs_data.channel == 2) + self.opt_data["cxff_pi.ti"] = cxrs_data.where(cxrs_data.channel == 2) if "efit" in self.diagnostics: - self.flat_data["efit.wp"] = ( + self.opt_data["efit.wp"] = ( self.models["efit"]() - .pop("wp") - .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) + .pop("wp") + .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) ) if noise: - self.flat_data["smmh1.ne"] = self.flat_data["smmh1.ne"] + self.flat_data[ + self.opt_data["smmh1.ne"] = self.opt_data["smmh1.ne"] + self.opt_data[ "smmh1.ne" ].max().values * np.random.normal(0, noise_factor, None) - self.flat_data["xrcs.spectra"] = self.flat_data[ - "xrcs.spectra" - ] + np.random.normal( + self.opt_data["xrcs.spectra"] = self.opt_data[ + "xrcs.spectra" + ] + np.random.normal( 0, - np.sqrt(self.flat_data["xrcs.spectra"].values[0,]), - self.flat_data["xrcs.spectra"].shape[1], + np.sqrt(self.opt_data["xrcs.spectra"].values[0,]), + self.opt_data["xrcs.spectra"].shape[1], ) - self.flat_data["cxff_pi.ti"] = self.flat_data[ - "cxff_pi.ti" - ] + self.flat_data["cxff_pi.ti"].max().values * np.random.normal( - 0, noise_factor, self.flat_data["cxff_pi.ti"].shape[1] + self.opt_data["cxff_pi.ti"] = self.opt_data[ + "cxff_pi.ti" + ] + self.opt_data["cxff_pi.ti"].max().values * np.random.normal( + 0, noise_factor, self.opt_data["cxff_pi.ti"].shape[1] ) - self.flat_data["efit.wp"] = self.flat_data["efit.wp"] + self.flat_data[ + self.opt_data["efit.wp"] = self.opt_data["efit.wp"] + self.opt_data[ "efit.wp" ].max().values * np.random.normal(0, noise_factor, None) - def exp_data(self): - self.flat_data = flatdict.FlatDict(self.ST40_data.binned_data, ".") + self.phantom_profiles = {} + for key in ["electron_density", "impurity_density", "electron_temperature", "ion_temperature", + "ion_density", "fast_density", "neutral_density"]: + self.phantom_profiles[key] = getattr(self.plasma, key).sel( + t=self.plasma.time_to_calculate + ).copy() + + def _exp_data(self): + self.opt_data = flatdict.FlatDict(self.data, ".") if "xrcs" in self.diagnostics: - self.flat_data["xrcs.spectra"]["wavelength"] = ( - self.flat_data["xrcs.spectra"].wavelength * 0.1 + self.opt_data["xrcs.spectra"]["wavelength"] = ( + self.opt_data["xrcs.spectra"].wavelength * 0.1 ) - background = self.flat_data["xrcs.spectra"].where( - (self.flat_data["xrcs.spectra"].wavelength < 0.392) - & (self.flat_data["xrcs.spectra"].wavelength > 0.388), + background = self.opt_data["xrcs.spectra"].where( + (self.opt_data["xrcs.spectra"].wavelength < 0.392) + & (self.opt_data["xrcs.spectra"].wavelength > 0.388), drop=True, ) - self.flat_data["xrcs.background"] = background.mean(dim="wavelength") - self.flat_data["xrcs.spectra"]["error"] = np.sqrt( - self.flat_data["xrcs.spectra"] + background.std(dim="wavelength") ** 2 + self.opt_data["xrcs.background"] = background.mean(dim="wavelength") + self.opt_data["xrcs.spectra"]["error"] = np.sqrt( + self.opt_data["xrcs.spectra"] + background.std(dim="wavelength") ** 2 ) if "cxff_pi" in self.diagnostics: - self.flat_data["cxff_pi"]["ti"] = self.flat_data["cxff_pi"]["ti"].where( - self.flat_data["cxff_pi"]["ti"].channel == 2, drop=True - ) - - def __call__(self, *args, **kwargs): - autocorr = sample_with_autocorr( - self.sampler, self.start_points, self.iterations, auto_sample=10 - ) - blobs = self.sampler.get_blobs( - discard=int(self.iterations * self.burn_in_fraction), flat=True - ) - blob_names = self.sampler.get_blobs().flatten()[0].keys() - blob_dict = { - blob_name: xr.concat( - [data[blob_name] for data in blobs], - dim=pd.Index(np.arange(0, blobs.__len__()), name="index"), + self.opt_data["cxff_pi"]["ti"] = self.opt_data["cxff_pi"]["ti"].where( + self.opt_data["cxff_pi"]["ti"].channel == 2, drop=True ) - for blob_name in blob_names - } - samples = self.sampler.get_chain(flat=True) - - prior_samples = self.bayes_run.sample_from_priors( - OPTIMISED_PARAMS, size=int(1e4) - ) - self.results = { - "blobs": blob_dict, - "diag_data": self.flat_data, - "samples": samples, - "prior_samples": prior_samples, - "param_names": OPTIMISED_PARAMS, - "phantom_profiles": self.profiles, - "autocorr": autocorr, - "acceptance_fraction": self.sampler.acceptance_fraction.sum(), - } - - self.nested_results = { - "DIAG_DATA":{"efit":{"wp":1}}, - # self.ST40_data.binned_data["efit"]["wp"].values[4,] - # "MODEL_DATA": {}, - "METADATA": { - "DIAG_RUNS": "EMPTY", - "EQUIL": self.ST40_data.equilibrium.__str__(), - "IMP1": self.plasma.impurities[0], - "MAIN_ION": self.plasma.main_ion, - }, - "OPTIMISATION": { - "AUTO_CORR":autocorr, - "BURN_FRAC":self.burn_in_fraction, - "ITER": self.iterations, - "NWALKERS": self.nwalkers, - "PARAM_NAMES": OPTIMISED_PARAMS, - "POST_SAMPLE": samples, - "PRIOR_SAMPLE": prior_samples - }, - # "PHANTOMS": self.profiles, - # "PROFILES": {}, - # "TIME": self.plasma.t, - # "TIME_OPT": self.plasma.time_to_calculate - - } - - # self.acceptance_fraction = self.sampler.acceptance_fraction.sum() - print(self.results["acceptance_fraction"]) - - Path(self.result_path).mkdir(parents=True, exist_ok=True) - with open(self.result_path + "results.pkl", "wb") as handle: - pickle.dump(self.results, handle) - - if self.save_plots: - plot_bayes_result(self.results, figheader=self.result_path, filetype=".png") - - -def sample_with_autocorr(sampler, start_points, iterations, auto_sample=5): - autocorr = np.ones((iterations,)) * np.nan - old_tau = np.inf - for sample in sampler.sample( - start_points, - iterations=iterations, - progress=True, - ): - if sampler.iteration % auto_sample: - continue - new_tau = np.mean(sampler.get_autocorr_time(tol=0)) - autocorr[sampler.iteration - 1] = new_tau - converged = np.all(new_tau * 50 < sampler.iteration) - converged &= np.all(np.abs(old_tau - new_tau) / new_tau < 0.01) - if converged: - break - old_tau = new_tau - autocorr = autocorr[: sampler.iteration, ] - return autocorr - + def __call__(self, filepath = "./results/test/", **kwargs): + self.run_sampler() + self.save_pickle(filepath=filepath) + plot_bayes_result(self.result, filepath) + return self.result if __name__ == "__main__": - mds_write=False - pulse_to_write = 23000101 - if mds_write: - node_structure = create_nodes(pulse_to_write=pulse_to_write, diagnostic_quantities=OPTIMISED_QUANTITY) + # mds_write=False + # pulse_to_write = 23000101 + # if mds_write: + # node_structure = create_nodes(pulse_to_write=pulse_to_write, diagnostic_quantities=OPTIMISED_QUANTITY) - run = BayesWorkflow( + run = DevBayesWorkflow( pulse=10009, - result_path="./results/test/", iterations=20, nwalkers=20, - burn_in_fraction=0.10, + burn_frac=0.10, dt=0.005, tsample=0.060, diagnostics=["efit", "smmh1", "cxff_pi"], - phantom=False, - center_mass_sampling=False, - Ti_ref=True, - mds_write=mds_write, - save_plots=True + phantoms=False, + sample_high_density=True, + mds_write=False, ) - run() + results = run(filepath="./results/test/",) - if mds_write: - write_nodes(pulse_to_write, node_structure, run.nested_results) + # if mds_write: + # write_nodes(pulse_to_write, node_structure, run.nested_results) From 633b36c7d13dc6c61a7b75c733784f92ee1837ff Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 4 Jul 2023 16:40:25 +0100 Subject: [PATCH 072/160] moved read data to abstract class --- indica/workflows/abstract_bayes_workflow.py | 19 ++++++++++--------- indica/workflows/bayes_workflow_dev.py | 8 -------- indica/workflows/bayes_workflow_example.py | 12 +----------- 3 files changed, 11 insertions(+), 28 deletions(-) diff --git a/indica/workflows/abstract_bayes_workflow.py b/indica/workflows/abstract_bayes_workflow.py index fae9a283..1faae4e6 100644 --- a/indica/workflows/abstract_bayes_workflow.py +++ b/indica/workflows/abstract_bayes_workflow.py @@ -5,6 +5,9 @@ from pathlib import Path import pickle +from indica.readers.read_st40 import ReadST40 + + class AbstractBayesWorkflow(ABC): @abstractmethod def __init__(self, @@ -35,16 +38,14 @@ def setup_plasma(self): """ self.plasma = None - @abstractmethod - def read_data(self, diagnostics: list): - """ - Reads data from server - - Returns - nested dictionary of data - """ - self.data = {} + def read_data(self, diagnostics: list): + self.reader = ReadST40( + self.pulse, tstart=self.tstart, tend=self.tend, dt=self.dt + ) + self.reader(diagnostics) + self.plasma.set_equilibrium(self.reader.equilibrium) + self.data = self.reader.binned_data @abstractmethod diff --git a/indica/workflows/bayes_workflow_dev.py b/indica/workflows/bayes_workflow_dev.py index ac6af031..28c228ed 100644 --- a/indica/workflows/bayes_workflow_dev.py +++ b/indica/workflows/bayes_workflow_dev.py @@ -193,14 +193,6 @@ def _init_fast_particles(self): Pbper = code_data["pbper"].interp(rho_poloidal=self.plasma.rho, t=self.plasma.t) self.plasma.pressure_fast_perpendicular.values = Pbper.values - def read_data(self, diagnostics: list): - self.reader = ReadST40( - self.pulse, tstart=self.tstart, tend=self.tend, dt=self.dt - ) - self.reader(diagnostics) - self.plasma.set_equilibrium(self.reader.equilibrium) - self.data = self.reader.binned_data - def setup_opt_data(self, phantoms=False): if phantoms: self._phantom_data() diff --git a/indica/workflows/bayes_workflow_example.py b/indica/workflows/bayes_workflow_example.py index b7677e18..1fd9d31e 100644 --- a/indica/workflows/bayes_workflow_example.py +++ b/indica/workflows/bayes_workflow_example.py @@ -171,15 +171,6 @@ def setup_plasma(self): self.plasma.update_profiles(DEFAULT_PHANTOM_PARAMS) self.plasma.build_atomic_data(calc_power_loss=False) - - def read_data(self, diagnostics: list): - self.reader = ReadST40( - self.pulse, tstart=self.tstart, tend=self.tend, dt=self.dt - ) - self.reader(diagnostics) - self.plasma.set_equilibrium(self.reader.equilibrium) - self.data = self.reader.binned_data - def setup_opt_data(self, phantoms=False): if phantoms: self._phantom_data() @@ -356,8 +347,7 @@ def __call__(self, filepath = "./results/test/", **kwargs): pulse=10009, dt=0.005, tsample=0.060, - # diagnostics=["efit", "smmh1", "cxff_pi"], - diagnostics=None, + diagnostics=["efit", "smmh1", "cxff_pi"], param_names=OPTIMISED_PARAMS, opt_quantity=OPTIMISED_QUANTITY, priors=DEFAULT_PRIORS, From 2ce1c5893724d5c4894f3900a046024da1b7e613 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Fri, 7 Jul 2023 10:07:46 +0100 Subject: [PATCH 073/160] initial commit --- indica/writers/bda_tree.py | 225 +++++++++++++++++++++++++++++++++++++ 1 file changed, 225 insertions(+) create mode 100644 indica/writers/bda_tree.py diff --git a/indica/writers/bda_tree.py b/indica/writers/bda_tree.py new file mode 100644 index 00000000..66126ee2 --- /dev/null +++ b/indica/writers/bda_tree.py @@ -0,0 +1,225 @@ +# Create trees from ppac standard utility tools +import standard_utility as util +from MDSplus import Connection +import sys +import numpy as np + + +def bda(): + nodes = { + "TIME": ("NUMERIC", "time vector, s"), + "TIME_OPT": ("NUMERIC", "time of optimisation, s"), + "METADATA": { + "GITCOMMIT": ("TEXT", "Commit ID used for run"), + "USER": ("TEXT", "Username of owner"), + "PULSE": ("NUMERIC", "Pulse number analysed"), + "EQUIL": ("TEXT", "Equilibrium used"), + "MAIN_ION": ("TEXT", "Main ion element"), + "IMP1": ("TEXT", "Impurity element chosen for Z1"), + "IMP2": ("TEXT", "Impurity element chosen for Z2"), + "DIAG_RUNS": ("TEXT", "RUNS from which diagnostic data originated") + }, + "PROFILES": { + "RHO_POLOIDAL": ("NUMERIC", "Radial vector, Sqrt of normalised poloidal flux"), + "NE": ("SIGNAL", "Electron density, m^-3"), + "NI": ("SIGNAL", "Ion density, m^-3"), + "TE": ("SIGNAL", "Electron temperature, eV"), + "TI": ("SIGNAL", "Ion temperature of main ion, eV"), + "NE_ERR": ("SIGNAL", "Electron density error, m^-3"), + "NI_ERR": ("SIGNAL", "Ion density error, m^-3"), + "TE_ERR": ("SIGNAL", "Electron temperature error, eV"), + "TI_ERR": ("SIGNAL", "Ion temperature of main ion error, eV"), + + + "TIZ1": ("SIGNAL", "Ion temperature of impurity IMP1, eV"), + "TIZ2": ("SIGNAL", "Ion temperature of impurity IMP2, eV"), + "TIZ3": ("SIGNAL", "Ion temperature of impurity IMP3, eV"), + "NIZ1": ("SIGNAL", "Density of impurity IMP1, m^-3"), + "NIZ2": ("SIGNAL", "Density of impurity IMP2, m^-3"), + "NIZ3": ("SIGNAL", "Density of impurity IMP3, m^-3"), + "NNEUTR": ("SIGNAL", "Density of neutral main ion, m^-3"), + "ZI": ("SIGNAL", "Average charge of main ion, "), + "ZIM1": ("SIGNAL", "Average charge of impurity IMP1, "), + "ZIM2": ("SIGNAL", "Average charge of impurity IMP2, "), + "ZIM3": ("SIGNAL", "Average charge of impurity IMP3, "), + "ZEFF": ("SIGNAL", "Effective charge, "), + "P": ("SIGNAL", "Pressure,Pa"), + "VOLUME": ("SIGNAL", "Volume inside magnetic surface,m^3"), + }, + + "PROFILE_STAT": { + "SAMPLES": ("NUMERIC", "Numerical index of the optimisation samples"), + "RHO_POLOIDAL": ("NUMERIC", "Radial vector, Sqrt of normalised poloidal flux"), + "NE": ("SIGNAL", "Electron density, m^-3"), + "NI": ("SIGNAL", "Ion density, m^-3"), + "TE": ("SIGNAL", "Electron temperature, eV"), + "TI": ("SIGNAL", "Ion temperature of main ion, eV"), + + "TIZ1": ("SIGNAL", "Ion temperature of impurity IMP1, eV"), + "TIZ2": ("SIGNAL", "Ion temperature of impurity IMP2, eV"), + "TIZ3": ("SIGNAL", "Ion temperature of impurity IMP3, eV"), + "NIZ1": ("SIGNAL", "Density of impurity IMP1, m^-3"), + "NIZ2": ("SIGNAL", "Density of impurity IMP2, m^-3"), + "NIZ3": ("SIGNAL", "Density of impurity IMP3, m^-3"), + "NNEUTR": ("SIGNAL", "Density of neutral main ion, m^-3"), + "ZI": ("SIGNAL", "Average charge of main ion, "), + "ZIM1": ("SIGNAL", "Average charge of impurity IMP1, "), + "ZIM2": ("SIGNAL", "Average charge of impurity IMP2, "), + "ZIM3": ("SIGNAL", "Average charge of impurity IMP3, "), + "ZEFF": ("SIGNAL", "Effective charge, "), + "P": ("SIGNAL", "Pressure,Pa"), + "VOLUME": ("SIGNAL", "Volume inside magnetic surface,m^3"), + }, + + "GLOBAL": { + "NE0": ("SIGNAL", "Central electron density, m^-3"), + "NI0": ("SIGNAL", "Central ion density, m^-3"), + "TE0": ("SIGNAL", "Central electron temperature, eV"), + "TI0": ("SIGNAL", "Central ion temperature of main ion, eV"), + + "NE0_ERR": ("SIGNAL", "Central electron density error, m^-3"), + "NI0_ERR": ("SIGNAL", "Central ion density error, m^-3"), + "TE0_ERR": ("SIGNAL", "Central electron temperature error, eV"), + "TI0_ERR": ("SIGNAL", "Central ion temperature of main ion error, eV"), + + }, + "PHANTOMS": { + "FLAG": ("TEXT", "True if phantoms used"), + "RHO_POLOIDAL": ("NUMERIC", "Radial vector, Sqrt of normalised poloidal flux"), + "NE": ("SIGNAL", "Electron density, m^-3"), + "NI": ("SIGNAL", "Ion density, m^-3"), + "TE": ("SIGNAL", "Electron temperature, eV"), + "TI": ("SIGNAL", "Ion temperature of main ion, eV"), + # "NIMP1": ("SIGNAL", "Impurity density, m^-3"), + + }, + + "INPUT": { + "PARAM_NAMES": ("TEXT", "Names of params optimised"), + "MODEL_KWARGS": ("TEXT", "Model key word arguments"), + "OPT_KWARGS": ("TEXT", "Optimiser key word arguments"), + "ITER": ("NUMERIC", "Iterations of optimiser"), + "NWALKERS": ("NUMERIC", "Number of walkers used"), + "BURN_FRAC": ("NUMERIC", "Burn in fraction for chains"), + }, + + "OPTIMISATION": { + "AUTO_CORR": ("NUMERIC", "Auto-correlation (iteration nwalker)"), + "POST_SAMPLE": ("NUMERIC", "Posterior probability samples (sample)"), + "PRIOR_SAMPLE": ("NUMERIC", "Prior samples"), + }, + } + return nodes + + +DIAGNOSTIC_QUANTITY = ["DIAGNOSTIC1.QUANTITY1", "DIAGNOSTIC1.QUANTITY2", "DIAGNOSTIC2.QUANTITY1"] + + +def create_nodes(pulse_to_write=23000101, run="RUN01", best=True, diagnostic_quantities=DIAGNOSTIC_QUANTITY): + bda_nodes = bda() + quant_list = [item.split(".") for item in diagnostic_quantities] # replace OPTIMISED_QUANTITY + diag_names = list(set([item[0] for item in quant_list])) + + + diag_nodes = {diag_name: + {quantity[1]: ("SIGNAL", f"measured {quantity[1]} from {quantity[0]}") + for quantity in quant_list if quantity[0] == diag_name} + for diag_name in diag_names + } + for diag_name in diag_names: + diag_nodes[diag_name]["RUN"] = ("TEXT", f"RUN from which {diag_name} data was taken") + + + nodes = { + "RUN": ("TEXT", "RUN used for diagnostic"), + "USAGE": ("TEXT", "Diagnostic used in analysis"), + "PULSE": ("NUMERIC", "Pulse used for diagnostic") + } + + workflow_nodes = {diag_name: nodes + for diag_name in diag_names + } + + + model_nodes = {diag_name: + {quantity[1]: ("SIGNAL", f"modelled {quantity[1]} from {quantity[0]}") + for quantity in quant_list if quantity[0] == diag_name} + for diag_name in diag_names + } + model_nodes["SAMPLES"] = ("NUMERIC", "index of samples taken from optimisation") + bda_nodes["MODEL_DATA"] = model_nodes + bda_nodes["DIAG_DATA"] = diag_nodes + bda_nodes["INPUT"]["WORKFLOW"] = workflow_nodes + + tree = "ST40" + script_name = "BDA" + script_info = "Bayesian Diagnostic Analysis" + + node_info = util.GetNodeInformation( + script=None, + node_information_type="json", + run_name=run, + run_info=run, + script_name=script_name, + script_info=script_info, + root_node=None, + tree=tree, + pulse_number=pulse_to_write, + base_node_to_read=None, + node_information_file=bda_nodes + ).get() + + util.StandardNodeCreation( + pulse_number=pulse_to_write, + dict_node_info=node_info, + mode="EDIT", + name_of_BEST="BEST", # name of the structure linked to BEST + link_BEST_to_run=best, + ) + return node_info + + +def write_nodes(pulse, node_info, data): + util.StandardNodeWriting( + pulse_number=pulse, # pulse number for which data should be written to the structure + dict_node_info=node_info, # node information file + nodes_to_write=[], # selective nodes to be written + data_to_write=data, + debug=False, + ) + + +def check_analysis_run(pulseNo, which_run, ): + # Checker function to see if data already exists in a run + IP_address_smaug = '192.168.1.7:8000' + conn = Connection(IP_address_smaug) + conn.openTree('BDA', pulseNo) + + temp = conn.get('BDA.' + which_run + '.TIME').data() + conn.closeAllTrees() + + overwrite_flag = True + if isinstance(temp, np.ndarray): + print(f'Data already Exists in pulseNo = {pulseNo}, which_run = {which_run}') + print('User prompt...') + question = f' Scheduled to overwrite pulseNo {pulseNo}, {which_run}' \ + f'\n Do you want to overwrite {which_run}? (y/n)' + overwrite_flag = query_yes_no(question) + return overwrite_flag + + +def query_yes_no(question, ): + valid = {"yes": True, "y": True, "no": False, "n": False} + while True: + sys.stdout.write(question) + choice = input().lower() + if choice == "": + return False + elif choice in valid.keys(): + return valid[choice] + else: + sys.stdout.write("Please respond with 'yes' or 'no' " "(or 'y' or 'n').\n") + + +if __name__ == "__main__": + create_nodes(pulse_to_write=23000101) From 08cf081a7fd05975de0c2e88f7ce3d0de4aac36d Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Fri, 7 Jul 2023 11:36:36 +0100 Subject: [PATCH 074/160] kwargs added --- indica/workflows/bayes_workflow_dev.py | 27 +++++++++++++++++--------- 1 file changed, 18 insertions(+), 9 deletions(-) diff --git a/indica/workflows/bayes_workflow_dev.py b/indica/workflows/bayes_workflow_dev.py index 28c228ed..d82eb87d 100644 --- a/indica/workflows/bayes_workflow_dev.py +++ b/indica/workflows/bayes_workflow_dev.py @@ -94,7 +94,11 @@ # "Ti_prof.peaking", ] -OPTIMISED_QUANTITY = ["cxff_pi.ti", "efit.wp", "smmh1.ne"] +OPTIMISED_QUANTITY = [ + # "xrcs.spectra", + "cxff_pi.ti", + "efit.wp", + "smmh1.ne"] class DevBayesWorkflow(AbstractBayesWorkflow): @@ -105,6 +109,7 @@ def __init__( param_names=None, opt_quantity=None, priors = None, + phantom_params=None, nwalkers=50, tstart=0.02, @@ -125,6 +130,7 @@ def __init__( self.param_names = param_names self.opt_quantity = opt_quantity self.priors = priors + self.phantom_params = phantom_params self.tstart = tstart self.tend = tend @@ -140,7 +146,7 @@ def __init__( self.mds_write = mds_write self.fast_particles = fast_particles - for attribute in ["pulse", "param_names", "opt_quantity", "priors", "diagnostics"]: + for attribute in ["pulse", "param_names", "opt_quantity", "priors", "diagnostics", "phantom_params"]: if getattr(self, attribute) is None: raise ValueError(f"{attribute} needs to be defined") @@ -168,7 +174,7 @@ def setup_plasma(self): self.plasma.time_to_calculate = self.plasma.t[ np.abs(self.tsample - self.plasma.t).argmin() ] - self.plasma.update_profiles(DEFAULT_PHANTOM_PARAMS) + self.plasma.update_profiles(self.phantom_params) self.plasma.build_atomic_data(calc_power_loss=False) if self.fast_particles: self._init_fast_particles() @@ -222,7 +228,6 @@ def setup_models(self, diagnostics: list): los_transform.set_equilibrium(self.plasma.equilibrium) model = Interferometry(name=diag) model.set_los_transform(los_transform) - model.plasma = self.plasma elif diag == "xrcs": los_transform = self.data["xrcs"]["te_kw"].transform @@ -235,20 +240,19 @@ def setup_models(self, diagnostics: list): * 0.1, ) model.set_los_transform(los_transform) - model.plasma = self.plasma elif diag == "efit": model = EquilibriumReconstruction(name="efit") - model.plasma = self.plasma elif diag == "cxff_pi": transform = self.data[diag]["ti"].transform transform.set_equilibrium(self.plasma.equilibrium) model = ChargeExchange(name=diag, element="ar") model.set_transect_transform(transform) - model.plasma = self.plasma else: raise ValueError(f"{diag} not found in setup_models") + model.plasma = self.plasma + self.models[diag] = model @@ -372,15 +376,20 @@ def __call__(self, filepath = "./results/test/", **kwargs): run = DevBayesWorkflow( pulse=10009, - iterations=20, + iterations=10, nwalkers=20, burn_frac=0.10, dt=0.005, tsample=0.060, - diagnostics=["efit", "smmh1", "cxff_pi"], + diagnostics=["xrcs", "efit", "smmh1", "cxff_pi"], phantoms=False, sample_high_density=True, mds_write=False, + + opt_quantity=OPTIMISED_QUANTITY, + param_names=OPTIMISED_PARAMS, + phantom_params=DEFAULT_PHANTOM_PARAMS, + priors=DEFAULT_PRIORS, ) results = run(filepath="./results/test/",) From 82eb00b40c27101f64a3d0597c0282b0bb8417c6 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Fri, 7 Jul 2023 14:24:25 +0100 Subject: [PATCH 075/160] xrcs wavelength units corrected --- indica/workflows/bayes_workflow_example.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/indica/workflows/bayes_workflow_example.py b/indica/workflows/bayes_workflow_example.py index 1fd9d31e..ea954787 100644 --- a/indica/workflows/bayes_workflow_example.py +++ b/indica/workflows/bayes_workflow_example.py @@ -210,7 +210,7 @@ def setup_models(self, diagnostics: list): window_vector=self.data[diag][ "spectra" ].wavelength.values - * 0.1, + , ) model.set_los_transform(los_transform) model.plasma = self.plasma From 666387bb7631fcccc014b1a5b680e84921a50c10 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Fri, 7 Jul 2023 14:32:25 +0100 Subject: [PATCH 076/160] background to int --- indica/models/helike_spectroscopy.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/indica/models/helike_spectroscopy.py b/indica/models/helike_spectroscopy.py index a0d3aade..3079ddd5 100644 --- a/indica/models/helike_spectroscopy.py +++ b/indica/models/helike_spectroscopy.py @@ -693,7 +693,7 @@ def __call__( if moment_analysis: self._calculate_temperatures() if background is not None: - self.measured_spectra = self.measured_spectra + background.sel(t=t) + self.measured_spectra = self.measured_spectra + background self._build_bckc_dictionary() return self.bckc From 3690efb9c15d054cca76cb6b59cdb49ce6b80a12 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 11 Jul 2023 11:13:50 +0100 Subject: [PATCH 077/160] stash --- indica/workflows/abstract_bayes_workflow.py | 82 ++++++++++++++------- indica/workflows/bayes_workflow_dev.py | 18 +++-- indica/writers/bda_tree.py | 33 +++++---- 3 files changed, 86 insertions(+), 47 deletions(-) diff --git a/indica/workflows/abstract_bayes_workflow.py b/indica/workflows/abstract_bayes_workflow.py index 1faae4e6..93265f00 100644 --- a/indica/workflows/abstract_bayes_workflow.py +++ b/indica/workflows/abstract_bayes_workflow.py @@ -64,7 +64,7 @@ def setup_models(self, diagnostics: list): self.models = {} @abstractmethod - def setup_optimiser(self): + def setup_optimiser(self, model_kwargs): """ Initialise and provide settings for optimiser """ @@ -93,32 +93,45 @@ def _build_result_dict(self): """ - result = {} result["TIME"] = self.plasma.t result["TIME_OPT"] = self.plasma.time_to_calculate + + quant_list = [item.split(".") for item in self.opt_quantity] + result["MODEL_DATA"] = {diag_name.upper(): + {quantity[1].upper(): self.blobs[f"{quantity[0]}.{quantity[1]}"] + for quantity in quant_list if quantity[0] == diag_name} + for diag_name in self.diagnostics + } + result["MODEL_DATA"]["SAMPLES"] = self.samples + + result["DIAG_DATA"] = {diag_name.upper(): + {quantity[1].upper(): self.opt_data[f"{quantity[0]}.{quantity[1]}"] + for quantity in quant_list if quantity[0] == diag_name} + for diag_name in self.diagnostics + } + result["INPUT"] = { "BURN_FRAC": self.burn_frac, "ITER": self.iterations, - "MODEL_KWARGS": "KWARGS", - "OPT_KWARGS": "KWARGS", + "MODEL_KWARGS": self.model_kwargs, "NWALKERS": self.nwalkers, "PARAM_NAMES": self.param_names, "PULSE": self.pulse, + "DT": self.dt, } - result["GLOBAL"] = { - "TI0": 0, - "TE0": 0, - "NE0": 0, - "NI0": 0, - "TI0_ERR": 0, - "TE0_ERR": 0, - "NE0_ERR": 0, - "NI0_ERR": 0, + result["INPUT"]["WORKFLOW"] = { + diag_name.upper(): + {"PULSE":self.pulse, # Change this if different pulses used for diagnostics + "USAGE": [quantity[1] for quantity in quant_list if quantity[0] == diag_name].join(), + "RUN": "PLACEHOLDER", + } + for diag_name in self.diagnostics } + result["PHANTOMS"] = { "FLAG": self.phantoms, "NE": self.phantom_profiles["electron_density"], @@ -131,6 +144,25 @@ def _build_result_dict(self): "NIMP2": self.phantom_profiles["impurity_density"].sel(element="c") } + result["PROFILES"] = { + "RHO_POLOIDAL": self.plasma.rho, + "RHO_TOR":self.plasma.equilibrium.rhotor, + + "NE": self.blobs["electron_density"].mean(dim="index"), + "NE_ERR": self.blobs["electron_density"].std(dim="index"), + + "NI": self.blobs["ion_density"].mean(dim="index"), + "TE": self.blobs["electron_temperature"].mean(dim="index"), + "TI": self.blobs["ion_temperature"].mean(dim="index"), + + + "NFAST": self.blobs["fast_density"], + "NNEUTR": self.blobs["neutral_density"], + "NIZ1": self.blobs["impurity_density"].sel(element="ar"), # TODO: generalise + "NIZ2": self.blobs["impurity_density"].sel(element="c") + } + + result["PROFILE_STAT"] = { "SAMPLES": self.samples, "RHO_POLOIDAL": self.plasma.rho, @@ -152,19 +184,19 @@ def _build_result_dict(self): "AUTOCORR": self.autocorr, } - quant_list = [item.split(".") for item in self.opt_quantity] - result["MODEL_DATA"] = {diag_name.upper(): - {quantity[1].upper(): self.blobs[f"{quantity[0]}.{quantity[1]}"] - for quantity in quant_list if quantity[0] == diag_name} - for diag_name in self.diagnostics - } - result["MODEL_DATA"]["SAMPLES"] = self.samples - result["DIAG_DATA"] = {diag_name.upper(): - {quantity[1].upper(): self.opt_data[f"{quantity[0]}.{quantity[1]}"] - for quantity in quant_list if quantity[0] == diag_name} - for diag_name in self.diagnostics - } + + result["GLOBAL"] = { + "TI0": 0, + "TE0": 0, + "NE0": 0, + "NI0": 0, + "TI0_ERR": 0, + "TE0_ERR": 0, + "NE0_ERR": 0, + "NI0_ERR": 0, + } + self.result = result return self.result diff --git a/indica/workflows/bayes_workflow_dev.py b/indica/workflows/bayes_workflow_dev.py index d82eb87d..b07e21b9 100644 --- a/indica/workflows/bayes_workflow_dev.py +++ b/indica/workflows/bayes_workflow_dev.py @@ -49,7 +49,7 @@ "Ne_prof.wcenter": get_uniform(0.1, 0.8), "Ne_prof.peaking": get_uniform(1, 6), "ar_conc": loguniform(0.0001, 0.01), - "Nimp_prof.y0": get_uniform(1e16, 1e18), + "Nimp_prof.y0": loguniform(1e16, 1e18), "Nimp_prof.y1": get_uniform(1e15, 2e16), "Ne_prof.y0/Nimp_prof.y0": lambda x1, x2: np.where( (x1 > x2 * 100) & (x1 < x2 * 1e5), 1, 0 @@ -95,7 +95,7 @@ ] OPTIMISED_QUANTITY = [ - # "xrcs.spectra", + "xrcs.spectra", "cxff_pi.ti", "efit.wp", "smmh1.ne"] @@ -106,10 +106,13 @@ def __init__( self, pulse=None, pulse_to_write=None, + diagnostics=None, param_names=None, opt_quantity=None, priors = None, + phantoms=False, phantom_params=None, + model_kwargs = None, nwalkers=50, tstart=0.02, @@ -119,10 +122,8 @@ def __init__( iterations=100, burn_frac=0, - phantoms=False, - diagnostics=None, sample_high_density = False, - mds_write=False, + mds_write = False, fast_particles = False, ): self.pulse = pulse @@ -131,6 +132,7 @@ def __init__( self.opt_quantity = opt_quantity self.priors = priors self.phantom_params = phantom_params + self.model_kwargs = model_kwargs self.tstart = tstart self.tend = tend @@ -237,7 +239,7 @@ def setup_models(self, diagnostics: list): window_vector=self.data[diag][ "spectra" ].wavelength.values - * 0.1, + , ) model.set_los_transform(los_transform) @@ -274,6 +276,7 @@ def setup_optimiser(self, ): log_prob_fn=self.bayesopt.ln_posterior, parameter_names=self.param_names, moves=self.move, + kwargs=self.model_kwargs, ) if self.sample_high_density: @@ -376,7 +379,7 @@ def __call__(self, filepath = "./results/test/", **kwargs): run = DevBayesWorkflow( pulse=10009, - iterations=10, + iterations=5, nwalkers=20, burn_frac=0.10, dt=0.005, @@ -390,6 +393,7 @@ def __call__(self, filepath = "./results/test/", **kwargs): param_names=OPTIMISED_PARAMS, phantom_params=DEFAULT_PHANTOM_PARAMS, priors=DEFAULT_PRIORS, + model_kwargs= {"moment_analysis":False, "background": 100} ) results = run(filepath="./results/test/",) diff --git a/indica/writers/bda_tree.py b/indica/writers/bda_tree.py index 66126ee2..1851f998 100644 --- a/indica/writers/bda_tree.py +++ b/indica/writers/bda_tree.py @@ -9,6 +9,21 @@ def bda(): nodes = { "TIME": ("NUMERIC", "time vector, s"), "TIME_OPT": ("NUMERIC", "time of optimisation, s"), + + "INPUT": { + "BURN_FRAC": ("NUMERIC", "Burn in fraction for chains"), + "ITER": ("NUMERIC", "Iterations of optimiser"), + "PARAM_NAMES": ("TEXT", "Names of params optimised"), + "MODEL_KWARGS": ("TEXT", "Model key word arguments"), + # "OPT_KWARGS": ("TEXT", "Optimiser key word arguments"), + "PULSE": ("NUMERIC", "Pulse number"), + + "TSTART": ("NUMERIC", "Start of time vector"), + "TEND": ("NUMERIC", "End of time vector"), + "DT": ("NUMERIC", "Distance between time points"), + "TSAMPLE": ("NUMERIC", "Sample time"), + }, + "METADATA": { "GITCOMMIT": ("TEXT", "Commit ID used for run"), "USER": ("TEXT", "Username of owner"), @@ -21,6 +36,7 @@ def bda(): }, "PROFILES": { "RHO_POLOIDAL": ("NUMERIC", "Radial vector, Sqrt of normalised poloidal flux"), + "RHO_TOR": ("NUMERIC", "Radial vector, toroidal flux"), "NE": ("SIGNAL", "Electron density, m^-3"), "NI": ("SIGNAL", "Ion density, m^-3"), "TE": ("SIGNAL", "Electron temperature, eV"), @@ -30,7 +46,6 @@ def bda(): "TE_ERR": ("SIGNAL", "Electron temperature error, eV"), "TI_ERR": ("SIGNAL", "Ion temperature of main ion error, eV"), - "TIZ1": ("SIGNAL", "Ion temperature of impurity IMP1, eV"), "TIZ2": ("SIGNAL", "Ion temperature of impurity IMP2, eV"), "TIZ3": ("SIGNAL", "Ion temperature of impurity IMP3, eV"), @@ -94,15 +109,6 @@ def bda(): }, - "INPUT": { - "PARAM_NAMES": ("TEXT", "Names of params optimised"), - "MODEL_KWARGS": ("TEXT", "Model key word arguments"), - "OPT_KWARGS": ("TEXT", "Optimiser key word arguments"), - "ITER": ("NUMERIC", "Iterations of optimiser"), - "NWALKERS": ("NUMERIC", "Number of walkers used"), - "BURN_FRAC": ("NUMERIC", "Burn in fraction for chains"), - }, - "OPTIMISATION": { "AUTO_CORR": ("NUMERIC", "Auto-correlation (iteration nwalker)"), "POST_SAMPLE": ("NUMERIC", "Posterior probability samples (sample)"), @@ -120,7 +126,6 @@ def create_nodes(pulse_to_write=23000101, run="RUN01", best=True, diagnostic_qua quant_list = [item.split(".") for item in diagnostic_quantities] # replace OPTIMISED_QUANTITY diag_names = list(set([item[0] for item in quant_list])) - diag_nodes = {diag_name: {quantity[1]: ("SIGNAL", f"measured {quantity[1]} from {quantity[0]}") for quantity in quant_list if quantity[0] == diag_name} @@ -129,7 +134,6 @@ def create_nodes(pulse_to_write=23000101, run="RUN01", best=True, diagnostic_qua for diag_name in diag_names: diag_nodes[diag_name]["RUN"] = ("TEXT", f"RUN from which {diag_name} data was taken") - nodes = { "RUN": ("TEXT", "RUN used for diagnostic"), "USAGE": ("TEXT", "Diagnostic used in analysis"), @@ -137,9 +141,8 @@ def create_nodes(pulse_to_write=23000101, run="RUN01", best=True, diagnostic_qua } workflow_nodes = {diag_name: nodes - for diag_name in diag_names - } - + for diag_name in diag_names + } model_nodes = {diag_name: {quantity[1]: ("SIGNAL", f"modelled {quantity[1]} from {quantity[0]}") From 3c23661233f5b196d1f7c1b6ab259bcf68e567c2 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 11 Jul 2023 13:52:15 +0100 Subject: [PATCH 078/160] renamed impurities --- indica/workflows/bayes_plots.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/indica/workflows/bayes_plots.py b/indica/workflows/bayes_plots.py index 9d781436..dc9ed77d 100644 --- a/indica/workflows/bayes_plots.py +++ b/indica/workflows/bayes_plots.py @@ -403,7 +403,7 @@ def plot_bayes_result( color="green", ) - key = "NIMP1" + key = "NIZ1" plot_profile( profiles[key], key, @@ -414,7 +414,7 @@ def plot_bayes_result( color="red", ) - key = "NIMP2" + key = "NIZ2" plot_profile( profiles[key], key, From bf3bfb9f1437fb8b9812d5f107c735fe5d65ff06 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Wed, 12 Jul 2023 11:58:24 +0100 Subject: [PATCH 079/160] stashing --- indica/workflows/abstract_bayes_workflow.py | 162 ++++++++++++++------ indica/workflows/bayes_workflow_dev.py | 64 ++++---- indica/writers/bda_tree.py | 98 ++++++++---- 3 files changed, 216 insertions(+), 108 deletions(-) diff --git a/indica/workflows/abstract_bayes_workflow.py b/indica/workflows/abstract_bayes_workflow.py index 93265f00..06f237b0 100644 --- a/indica/workflows/abstract_bayes_workflow.py +++ b/indica/workflows/abstract_bayes_workflow.py @@ -88,114 +88,164 @@ def save_phantom_profiles(self, kinetic_profiles=None): def _build_result_dict(self): """ - TODO: xarray to numpy for MDSPlus writing (plotting / dimensions / units have to be figured out) + Returns + ------- + + dictionary of results in MDS+ structure """ result = {} + quant_list = [item.split(".") for item in self.opt_quantity] + result["TIME"] = self.plasma.t result["TIME_OPT"] = self.plasma.time_to_calculate - - quant_list = [item.split(".") for item in self.opt_quantity] - result["MODEL_DATA"] = {diag_name.upper(): - {quantity[1].upper(): self.blobs[f"{quantity[0]}.{quantity[1]}"] - for quantity in quant_list if quantity[0] == diag_name} - for diag_name in self.diagnostics - } - result["MODEL_DATA"]["SAMPLES"] = self.samples - - result["DIAG_DATA"] = {diag_name.upper(): - {quantity[1].upper(): self.opt_data[f"{quantity[0]}.{quantity[1]}"] - for quantity in quant_list if quantity[0] == diag_name} - for diag_name in self.diagnostics - } + result["METADATA"] = { + "GITCOMMIT": "PLACEHOLDER", + "USER": "PLACEHOLDER", + "EQUIL": "PLACEHOLDER", + } result["INPUT"] = { "BURN_FRAC": self.burn_frac, "ITER": self.iterations, - "MODEL_KWARGS": self.model_kwargs, "NWALKERS": self.nwalkers, + + "MODEL_KWARGS": self.model_kwargs, + "OPT_QUANTITY": self.opt_quantity, "PARAM_NAMES": self.param_names, + "PULSE": self.pulse, "DT": self.dt, + "IMPURITIES":self.plasma.impurities, + "MAIN_ION": self.plasma.main_ion, } - result["INPUT"]["WORKFLOW"] = { diag_name.upper(): - {"PULSE":self.pulse, # Change this if different pulses used for diagnostics - "USAGE": [quantity[1] for quantity in quant_list if quantity[0] == diag_name].join(), + {"PULSE": self.pulse, # Change this if different pulses used for diagnostics + "USAGE": "".join([quantity[1] for quantity in quant_list if quantity[0] == diag_name]), "RUN": "PLACEHOLDER", } for diag_name in self.diagnostics } + result["MODEL_DATA"] = {diag_name.upper(): + {quantity[1].upper(): self.blobs[f"{quantity[0]}.{quantity[1]}"] + for quantity in quant_list if quantity[0] == diag_name} + for diag_name in self.diagnostics + } + result["MODEL_DATA"]["SAMPLES"] = self.samples + + result["DIAG_DATA"] = {diag_name.upper(): + {quantity[1].upper(): self.opt_data[f"{quantity[0]}.{quantity[1]}"] + for quantity in quant_list if quantity[0] == diag_name} + for diag_name in self.diagnostics + } + result["PHANTOMS"] = { "FLAG": self.phantoms, "NE": self.phantom_profiles["electron_density"], "TE": self.phantom_profiles["electron_temperature"], - "TI": self.phantom_profiles["ion_temperature"], - "NI": self.phantom_profiles["ion_density"], + "TI": self.phantom_profiles["ion_temperature"].sel(element=self.plasma.main_ion), + "NI": self.phantom_profiles["ion_density"].sel(element=self.plasma.main_ion), "NNEUTR": self.phantom_profiles["neutral_density"], "NFAST": self.phantom_profiles["fast_density"], - "NIMP1": self.phantom_profiles["impurity_density"].sel(element="ar"), # TODO: generalise - "NIMP2": self.phantom_profiles["impurity_density"].sel(element="c") } + result["PHANTOMS"] = {**result["PHANTOMS"], **{ + f"NIZ{num_imp+1}": self.phantom_profiles["impurity_density"].sel(element=imp) + for num_imp, imp in enumerate(self.plasma.impurities)}} + result["PHANTOMS"] = {**result["PHANTOMS"], **{ + f"TIZ{num_imp + 1}": self.phantom_profiles["ion_temperature"].sel(element=imp) + for num_imp, imp in enumerate(self.plasma.impurities)}} + result["PROFILES"] = { "RHO_POLOIDAL": self.plasma.rho, - "RHO_TOR":self.plasma.equilibrium.rhotor, - - "NE": self.blobs["electron_density"].mean(dim="index"), - "NE_ERR": self.blobs["electron_density"].std(dim="index"), + "RHO_TOR": self.plasma.equilibrium.rhotor.interp(t=self.plasma.t), - "NI": self.blobs["ion_density"].mean(dim="index"), - "TE": self.blobs["electron_temperature"].mean(dim="index"), - "TI": self.blobs["ion_temperature"].mean(dim="index"), + "NE": self.blobs["electron_density"].median(dim="index"), + "NI": self.blobs["ion_density"].sel(element=self.plasma.main_ion).median(dim="index"), + "TE": self.blobs["electron_temperature"].median(dim="index"), + "TI": self.blobs["ion_temperature"].sel(element=self.plasma.main_ion).median(dim="index"), + "NFAST": self.blobs["fast_density"].median(dim="index"), + "NNEUTR": self.blobs["neutral_density"].median(dim="index"), + "NE_ERR": self.blobs["electron_density"].std(dim="index"), + "NI_ERR": self.blobs["ion_density"].sel(element=self.plasma.main_ion).std(dim="index"), + "TE_ERR": self.blobs["electron_temperature"].std(dim="index"), + "TI_ERR": self.blobs["ion_temperature"].sel(element=self.plasma.main_ion).std(dim="index"), + "NFAST_ERR": self.blobs["fast_density"].std(dim="index"), + "NNEUTR_ERR": self.blobs["neutral_density"].std(dim="index"), - "NFAST": self.blobs["fast_density"], - "NNEUTR": self.blobs["neutral_density"], - "NIZ1": self.blobs["impurity_density"].sel(element="ar"), # TODO: generalise - "NIZ2": self.blobs["impurity_density"].sel(element="c") } + result["PROFILES"] = {**result["PROFILES"], **{ + f"NIZ{num_imp+1}": self.blobs["impurity_density"].sel(element=imp).median(dim="index") + for num_imp, imp in enumerate(self.plasma.impurities)}} + result["PROFILES"] = {**result["PROFILES"], **{ + f"NIZ{num_imp+1}_ERR": self.blobs["impurity_density"].sel(element=imp).std(dim="index") + for num_imp, imp in enumerate(self.plasma.impurities)}} + result["PROFILES"] = {**result["PROFILES"], **{ + f"TIZ{num_imp + 1}": self.blobs["ion_temperature"].sel(element=imp).median(dim="index") + for num_imp, imp in enumerate(self.plasma.impurities)}} + result["PROFILES"] = {**result["PROFILES"], **{ + f"TIZ{num_imp + 1}_ERR": self.blobs["ion_temperature"].sel(element=imp).std(dim="index") + for num_imp, imp in enumerate(self.plasma.impurities)}} result["PROFILE_STAT"] = { "SAMPLES": self.samples, "RHO_POLOIDAL": self.plasma.rho, "NE": self.blobs["electron_density"], - "NI": self.blobs["ion_density"], + "NI": self.blobs["ion_density"].sel(element=self.plasma.main_ion), "TE": self.blobs["electron_temperature"], - "TI": self.blobs["ion_temperature"], + "TI": self.blobs["ion_temperature"].sel(element=self.plasma.main_ion), "NFAST": self.blobs["fast_density"], "NNEUTR": self.blobs["neutral_density"], - "NIMP1": self.blobs["impurity_density"].sel(element="ar"), # TODO: generalise - "NIMP2": self.blobs["impurity_density"].sel(element="c") } + result["PROFILE_STAT"] = {**result["PROFILE_STAT"], **{ + f"NIZ{num_imp+1}": self.blobs["impurity_density"].sel(element=imp) + for num_imp, imp in enumerate(self.plasma.impurities)}} + result["PROFILE_STAT"] = {**result["PROFILE_STAT"], **{ + f"TIZ{num_imp + 1}": self.blobs["ion_temperature"].sel(element=imp) + for num_imp, imp in enumerate(self.plasma.impurities)}} + result["OPTIMISATION"] = { - "OPT_QUANTITY": self.opt_quantity, "ACCEPT_FRAC": self.accept_frac, "PRIOR_SAMPLE": self.prior_sample, "POST_SAMPLE": self.post_sample, "AUTOCORR": self.autocorr, } - - result["GLOBAL"] = { - "TI0": 0, - "TE0": 0, - "NE0": 0, - "NI0": 0, - "TI0_ERR": 0, - "TE0_ERR": 0, - "NE0_ERR": 0, - "NI0_ERR": 0, + "TI0": self.blobs["ion_temperature"].sel(element=self.plasma.main_ion).sel(rho_poloidal=0, method="nearest").median(dim="index"), + "TE0": self.blobs["electron_temperature"].sel(rho_poloidal=0, method="nearest").median(dim="index"), + "NE0": self.blobs["electron_density"].sel(rho_poloidal=0, method="nearest").median(dim="index"), + "NI0": self.blobs["ion_density"].sel(element=self.plasma.main_ion).sel(rho_poloidal=0, method="nearest").median(dim="index"), + "TI0_ERR": self.blobs["ion_temperature"].sel(element=self.plasma.main_ion).sel(rho_poloidal=0, method="nearest").std(dim="index"), + "TE0_ERR": self.blobs["electron_temperature"].sel(rho_poloidal=0, method="nearest").std(dim="index"), + "NE0_ERR": self.blobs["electron_density"].sel(rho_poloidal=0, method="nearest").std(dim="index"), + "NI0_ERR": self.blobs["ion_density"].sel(element=self.plasma.main_ion).sel(rho_poloidal=0, method="nearest").std(dim="index"), } + result["GLOBAL"] = {**result["GLOBAL"], **{ + f"TI0Z{num_imp + 1}": self.blobs["ion_temperature"].sel(element=imp).sel(rho_poloidal=0, method="nearest").median(dim="index") + for num_imp, imp in enumerate(self.plasma.impurities)}} + result["GLOBAL"] = {**result["GLOBAL"], **{ + f"TI0Z{num_imp + 1}_ERR": self.blobs["ion_temperature"].sel(element=imp).sel(rho_poloidal=0, + method="nearest").std( + dim="index") + for num_imp, imp in enumerate(self.plasma.impurities)}} + result["GLOBAL"] = {**result["GLOBAL"], **{ + f"NI0Z{num_imp + 1}": self.blobs["impurity_density"].sel(element=imp).sel(rho_poloidal=0, method="nearest").median(dim="index") + for num_imp, imp in enumerate(self.plasma.impurities)}} + result["GLOBAL"] = {**result["GLOBAL"], **{ + f"NI0Z{num_imp + 1}_ERR": self.blobs["impurity_density"].sel(element=imp).sel(rho_poloidal=0, + method="nearest").std( + dim="index") + for num_imp, imp in enumerate(self.plasma.impurities)}} self.result = result return self.result @@ -242,10 +292,24 @@ def save_pickle(self, filepath): with open(filepath + "results.pkl", "wb") as handle: pickle.dump(self.result, handle) + def dict_of_dataarray_to_numpy(self, dict_of_dataarray): + """ + Mutates input dictionary to change xr.DataArray objects to np.array + + """ + for key, value in dict_of_dataarray.items(): + if isinstance(value, dict): + self.dict_of_dataarray_to_numpy(value) + elif isinstance(value, xr.DataArray): + dict_of_dataarray[key] = dict_of_dataarray[key].values + return dict_of_dataarray + @abstractmethod def __call__(self, filepath="./results/test/", **kwargs): self.run_sampler() self.save_pickle(filepath=filepath) + self.result = self.dict_of_dataarray_to_numpy(self.result) + return self.result diff --git a/indica/workflows/bayes_workflow_dev.py b/indica/workflows/bayes_workflow_dev.py index b07e21b9..ba0e97ed 100644 --- a/indica/workflows/bayes_workflow_dev.py +++ b/indica/workflows/bayes_workflow_dev.py @@ -1,6 +1,7 @@ import emcee import numpy as np import flatdict +import copy from scipy.stats import loguniform from indica.readers.read_st40 import ReadST40 @@ -122,17 +123,20 @@ def __init__( iterations=100, burn_frac=0, + mds_write=False, + plot=True, sample_high_density = False, - mds_write = False, fast_particles = False, ): self.pulse = pulse self.pulse_to_write = pulse_to_write + self.diagnostics = diagnostics self.param_names = param_names self.opt_quantity = opt_quantity self.priors = priors self.phantom_params = phantom_params self.model_kwargs = model_kwargs + self.phantoms = phantoms self.tstart = tstart self.tend = tend @@ -142,10 +146,9 @@ def __init__( self.iterations = iterations self.burn_frac = burn_frac - self.phantoms = phantoms - self.diagnostics = diagnostics - self.sample_high_density = sample_high_density self.mds_write = mds_write + self.plot = plot + self.sample_high_density = sample_high_density self.fast_particles = fast_particles for attribute in ["pulse", "param_names", "opt_quantity", "priors", "diagnostics", "phantom_params"]: @@ -154,10 +157,10 @@ def __init__( self.setup_plasma() self.save_phantom_profiles() - self.read_data(diagnostics) + self.read_data(self.diagnostics) self.setup_opt_data(phantoms=self.phantoms) - self.setup_models(diagnostics) - self.setup_optimiser() + self.setup_models(self.diagnostics) + self.setup_optimiser(self.model_kwargs) def setup_plasma(self): self.plasma = Plasma( @@ -258,7 +261,7 @@ def setup_models(self, diagnostics: list): self.models[diag] = model - def setup_optimiser(self, ): + def setup_optimiser(self, model_kwargs): self.bayesopt = BayesModels( plasma=self.plasma, @@ -276,7 +279,7 @@ def setup_optimiser(self, ): log_prob_fn=self.bayesopt.ln_posterior, parameter_names=self.param_names, moves=self.move, - kwargs=self.model_kwargs, + kwargs=model_kwargs, ) if self.sample_high_density: @@ -364,38 +367,47 @@ def _exp_data(self): ) def __call__(self, filepath = "./results/test/", **kwargs): + + if self.mds_write: + self.node_structure = create_nodes(pulse_to_write=self.pulse_to_write, + diagnostic_quantities=self.opt_quantity, + mode="NEW") + self.run_sampler() self.save_pickle(filepath=filepath) - plot_bayes_result(self.result, filepath) + self.result = self.dict_of_dataarray_to_numpy(self.result) + + if self.plot: + plot_bayes_result(self.result, filepath) + + if self.mds_write: + write_nodes(self.pulse_to_write, self.node_structure, self.result) + return self.result if __name__ == "__main__": - # mds_write=False - # pulse_to_write = 23000101 - # if mds_write: - # node_structure = create_nodes(pulse_to_write=pulse_to_write, diagnostic_quantities=OPTIMISED_QUANTITY) - run = DevBayesWorkflow( pulse=10009, + pulse_to_write=23000101, + diagnostics=["xrcs", "efit", "smmh1", "cxff_pi"], + opt_quantity=OPTIMISED_QUANTITY, + param_names=OPTIMISED_PARAMS, + phantom_params=DEFAULT_PHANTOM_PARAMS, + priors=DEFAULT_PRIORS, + iterations=5, nwalkers=20, burn_frac=0.10, dt=0.005, tsample=0.060, - diagnostics=["xrcs", "efit", "smmh1", "cxff_pi"], - phantoms=False, - sample_high_density=True, - mds_write=False, - opt_quantity=OPTIMISED_QUANTITY, - param_names=OPTIMISED_PARAMS, - phantom_params=DEFAULT_PHANTOM_PARAMS, - priors=DEFAULT_PRIORS, - model_kwargs= {"moment_analysis":False, "background": 100} + mds_write=True, + plot=False, + phantoms=False, + sample_high_density=False, + model_kwargs= { "background": 100} ) results = run(filepath="./results/test/",) - # if mds_write: - # write_nodes(pulse_to_write, node_structure, run.nested_results) diff --git a/indica/writers/bda_tree.py b/indica/writers/bda_tree.py index 1851f998..1d16083c 100644 --- a/indica/writers/bda_tree.py +++ b/indica/writers/bda_tree.py @@ -1,4 +1,5 @@ # Create trees from ppac standard utility tools + import standard_utility as util from MDSplus import Connection import sys @@ -13,46 +14,44 @@ def bda(): "INPUT": { "BURN_FRAC": ("NUMERIC", "Burn in fraction for chains"), "ITER": ("NUMERIC", "Iterations of optimiser"), - "PARAM_NAMES": ("TEXT", "Names of params optimised"), + "PARAM_NAMES": ("TEXT", "Names of parameters optimised"), + "OPT_QUANTITY":("TEXT", "Names of quantities optimised"), "MODEL_KWARGS": ("TEXT", "Model key word arguments"), # "OPT_KWARGS": ("TEXT", "Optimiser key word arguments"), "PULSE": ("NUMERIC", "Pulse number"), - "TSTART": ("NUMERIC", "Start of time vector"), - "TEND": ("NUMERIC", "End of time vector"), - "DT": ("NUMERIC", "Distance between time points"), - "TSAMPLE": ("NUMERIC", "Sample time"), + "TSTART": ("NUMERIC", "Start of time vector, s"), + "TEND": ("NUMERIC", "End of time vector, s"), + "DT": ("NUMERIC", "Distance between time points, s"), + "TSAMPLE": ("NUMERIC", "Sample time, s"), + "IMPURITIES": ("TEXT", "Names of impurity elements"), + "MAIN_ION": ("TEXT", "Name of main ion"), + }, "METADATA": { "GITCOMMIT": ("TEXT", "Commit ID used for run"), "USER": ("TEXT", "Username of owner"), - "PULSE": ("NUMERIC", "Pulse number analysed"), "EQUIL": ("TEXT", "Equilibrium used"), - "MAIN_ION": ("TEXT", "Main ion element"), - "IMP1": ("TEXT", "Impurity element chosen for Z1"), - "IMP2": ("TEXT", "Impurity element chosen for Z2"), - "DIAG_RUNS": ("TEXT", "RUNS from which diagnostic data originated") }, + "PROFILES": { - "RHO_POLOIDAL": ("NUMERIC", "Radial vector, Sqrt of normalised poloidal flux"), - "RHO_TOR": ("NUMERIC", "Radial vector, toroidal flux"), + "RHOP": ("NUMERIC", "Radial vector, Sqrt of normalised poloidal flux"), + "RHOT": ("SIGNAL", "Radial vector, toroidal flux"), + "NE": ("SIGNAL", "Electron density, m^-3"), "NI": ("SIGNAL", "Ion density, m^-3"), "TE": ("SIGNAL", "Electron temperature, eV"), "TI": ("SIGNAL", "Ion temperature of main ion, eV"), - "NE_ERR": ("SIGNAL", "Electron density error, m^-3"), - "NI_ERR": ("SIGNAL", "Ion density error, m^-3"), - "TE_ERR": ("SIGNAL", "Electron temperature error, eV"), - "TI_ERR": ("SIGNAL", "Ion temperature of main ion error, eV"), - - "TIZ1": ("SIGNAL", "Ion temperature of impurity IMP1, eV"), - "TIZ2": ("SIGNAL", "Ion temperature of impurity IMP2, eV"), - "TIZ3": ("SIGNAL", "Ion temperature of impurity IMP3, eV"), - "NIZ1": ("SIGNAL", "Density of impurity IMP1, m^-3"), - "NIZ2": ("SIGNAL", "Density of impurity IMP2, m^-3"), - "NIZ3": ("SIGNAL", "Density of impurity IMP3, m^-3"), + "TIZ1": ("SIGNAL", "Ion temperature of impurity Z1, eV"), + "TIZ2": ("SIGNAL", "Ion temperature of impurity Z2, eV"), + "TIZ3": ("SIGNAL", "Ion temperature of impurity Z3, eV"), + "NIZ1": ("SIGNAL", "Density of impurity Z1, m^-3"), + "NIZ2": ("SIGNAL", "Density of impurity Z2, m^-3"), + "NIZ3": ("SIGNAL", "Density of impurity Z3, m^-3"), "NNEUTR": ("SIGNAL", "Density of neutral main ion, m^-3"), + "NFAST": ("SIGNAL", "Density of fast ion, m^-3"), + "ZI": ("SIGNAL", "Average charge of main ion, "), "ZIM1": ("SIGNAL", "Average charge of impurity IMP1, "), "ZIM2": ("SIGNAL", "Average charge of impurity IMP2, "), @@ -60,16 +59,29 @@ def bda(): "ZEFF": ("SIGNAL", "Effective charge, "), "P": ("SIGNAL", "Pressure,Pa"), "VOLUME": ("SIGNAL", "Volume inside magnetic surface,m^3"), + + "NE_ERR": ("SIGNAL", "Electron density error, m^-3"), + "NI_ERR": ("SIGNAL", "Ion density error, m^-3"), + "TE_ERR": ("SIGNAL", "Electron temperature error, eV"), + "TI_ERR": ("SIGNAL", "Ion temperature of main ion error, eV"), + "TIZ1_ERR": ("SIGNAL", "Ion temperature of impurity Z1 error, eV"), + "TIZ2_ERR": ("SIGNAL", "Ion temperature of impurity Z2 error, eV"), + "TIZ3_ERR": ("SIGNAL", "Ion temperature of impurity Z3 error, eV"), + "NIZ1_ERR": ("SIGNAL", "Density of impurity Z1 error, m^-3"), + "NIZ2_ERR": ("SIGNAL", "Density of impurity Z2 error, m^-3"), + "NIZ3_ERR": ("SIGNAL", "Density of impurity Z3 error, m^-3"), + "NNEUTR_ERR": ("SIGNAL", "Density of neutral main ion error, m^-3"), + "NFAST_ERR": ("SIGNAL", "Density of fast ion error, m^-3"), + }, "PROFILE_STAT": { "SAMPLES": ("NUMERIC", "Numerical index of the optimisation samples"), - "RHO_POLOIDAL": ("NUMERIC", "Radial vector, Sqrt of normalised poloidal flux"), + "RHOP": ("NUMERIC", "Radial vector, Sqrt of normalised poloidal flux"), "NE": ("SIGNAL", "Electron density, m^-3"), "NI": ("SIGNAL", "Ion density, m^-3"), "TE": ("SIGNAL", "Electron temperature, eV"), "TI": ("SIGNAL", "Ion temperature of main ion, eV"), - "TIZ1": ("SIGNAL", "Ion temperature of impurity IMP1, eV"), "TIZ2": ("SIGNAL", "Ion temperature of impurity IMP2, eV"), "TIZ3": ("SIGNAL", "Ion temperature of impurity IMP3, eV"), @@ -77,13 +89,15 @@ def bda(): "NIZ2": ("SIGNAL", "Density of impurity IMP2, m^-3"), "NIZ3": ("SIGNAL", "Density of impurity IMP3, m^-3"), "NNEUTR": ("SIGNAL", "Density of neutral main ion, m^-3"), + "NFAST": ("SIGNAL", "Density of fast ions, m^-3"), + "ZI": ("SIGNAL", "Average charge of main ion, "), "ZIM1": ("SIGNAL", "Average charge of impurity IMP1, "), "ZIM2": ("SIGNAL", "Average charge of impurity IMP2, "), "ZIM3": ("SIGNAL", "Average charge of impurity IMP3, "), "ZEFF": ("SIGNAL", "Effective charge, "), "P": ("SIGNAL", "Pressure,Pa"), - "VOLUME": ("SIGNAL", "Volume inside magnetic surface,m^3"), + "VOLUME": ("SIGNAL", "Volume inside magnetic surface, m^3"), }, "GLOBAL": { @@ -91,11 +105,23 @@ def bda(): "NI0": ("SIGNAL", "Central ion density, m^-3"), "TE0": ("SIGNAL", "Central electron temperature, eV"), "TI0": ("SIGNAL", "Central ion temperature of main ion, eV"), + "TI0Z1": ("SIGNAL", "Central ion temperature of impurity Z1, eV"), + "TI0Z2": ("SIGNAL", "Central ion temperature of impurity Z2, eV"), + "TI0Z3": ("SIGNAL", "Central ion temperature of impurity Z3, eV"), + "NI0Z1": ("SIGNAL", "Central density of impurity Z1, m^-3"), + "NI0Z2": ("SIGNAL", "Central density of impurity Z2, m^-3"), + "NI0Z3": ("SIGNAL", "Central density of impurity Z3, m^-3"), "NE0_ERR": ("SIGNAL", "Central electron density error, m^-3"), "NI0_ERR": ("SIGNAL", "Central ion density error, m^-3"), "TE0_ERR": ("SIGNAL", "Central electron temperature error, eV"), "TI0_ERR": ("SIGNAL", "Central ion temperature of main ion error, eV"), + "TI0Z1_ERR": ("SIGNAL", "Central ion temperature of impurity Z1, eV"), + "TI0Z2_ERR": ("SIGNAL", "Central ion temperature of impurity Z2, eV"), + "TI0Z3_ERR": ("SIGNAL", "Central ion temperature of impurity Z3, eV"), + "NI0Z1_ERR": ("SIGNAL", "Central density of impurity Z1, m^-3"), + "NI0Z2_ERR": ("SIGNAL", "Central density of impurity Z2, m^-3"), + "NI0Z3_ERR": ("SIGNAL", "Central density of impurity Z3, m^-3"), }, "PHANTOMS": { @@ -105,11 +131,17 @@ def bda(): "NI": ("SIGNAL", "Ion density, m^-3"), "TE": ("SIGNAL", "Electron temperature, eV"), "TI": ("SIGNAL", "Ion temperature of main ion, eV"), - # "NIMP1": ("SIGNAL", "Impurity density, m^-3"), + "TIZ1": ("SIGNAL", "Ion temperature of impurity Z1 , eV"), + "TIZ2": ("SIGNAL", "Ion temperature of impurity Z2, eV"), + "TIZ3": ("SIGNAL", "Ion temperature of impurity Z3, eV"), + "NIZ1": ("SIGNAL", "Impurity density of Z1, m^-3"), + "NIZ2": ("SIGNAL", "Impurity density of Z2, m^-3"), + "NIZ3": ("SIGNAL", "Impurity density of Z3, m^-3"), }, "OPTIMISATION": { + "ACCEPT_FRAC": ("NUMERIC", "Fraction of samples accepted by optimiser"), "AUTO_CORR": ("NUMERIC", "Auto-correlation (iteration nwalker)"), "POST_SAMPLE": ("NUMERIC", "Posterior probability samples (sample)"), "PRIOR_SAMPLE": ("NUMERIC", "Prior samples"), @@ -121,9 +153,9 @@ def bda(): DIAGNOSTIC_QUANTITY = ["DIAGNOSTIC1.QUANTITY1", "DIAGNOSTIC1.QUANTITY2", "DIAGNOSTIC2.QUANTITY1"] -def create_nodes(pulse_to_write=23000101, run="RUN01", best=True, diagnostic_quantities=DIAGNOSTIC_QUANTITY): +def create_nodes(pulse_to_write=23000101, run="RUN01", best=True, diagnostic_quantities=DIAGNOSTIC_QUANTITY, mode="EDIT"): bda_nodes = bda() - quant_list = [item.split(".") for item in diagnostic_quantities] # replace OPTIMISED_QUANTITY + quant_list = [item.upper().split(".") for item in diagnostic_quantities] # replace OPTIMISED_QUANTITY diag_names = list(set([item[0] for item in quant_list])) diag_nodes = {diag_name: @@ -131,12 +163,12 @@ def create_nodes(pulse_to_write=23000101, run="RUN01", best=True, diagnostic_qua for quantity in quant_list if quantity[0] == diag_name} for diag_name in diag_names } - for diag_name in diag_names: - diag_nodes[diag_name]["RUN"] = ("TEXT", f"RUN from which {diag_name} data was taken") + # for diag_name in diag_names: + # diag_nodes[diag_name]["RUN"] = ("TEXT", f"RUN from which {diag_name} data was taken") nodes = { "RUN": ("TEXT", "RUN used for diagnostic"), - "USAGE": ("TEXT", "Diagnostic used in analysis"), + "USAGE": ("TEXT", "Quantity used in analysis"), "PULSE": ("NUMERIC", "Pulse used for diagnostic") } @@ -175,7 +207,7 @@ def create_nodes(pulse_to_write=23000101, run="RUN01", best=True, diagnostic_qua util.StandardNodeCreation( pulse_number=pulse_to_write, dict_node_info=node_info, - mode="EDIT", + mode=mode, name_of_BEST="BEST", # name of the structure linked to BEST link_BEST_to_run=best, ) From fbfafe5523649bf4a380bb4750a8b0e63ab38aca Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 13 Jul 2023 15:17:41 +0100 Subject: [PATCH 080/160] stashing --- indica/workflows/abstract_bayes_workflow.py | 11 ++++++----- indica/workflows/bayes_workflow_dev.py | 6 +++++- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/indica/workflows/abstract_bayes_workflow.py b/indica/workflows/abstract_bayes_workflow.py index 06f237b0..86d1fb99 100644 --- a/indica/workflows/abstract_bayes_workflow.py +++ b/indica/workflows/abstract_bayes_workflow.py @@ -154,13 +154,14 @@ def _build_result_dict(self): "NNEUTR": self.phantom_profiles["neutral_density"], "NFAST": self.phantom_profiles["fast_density"], } - result["PHANTOMS"] = {**result["PHANTOMS"], **{ + result["PHANTOMS"].update({ f"NIZ{num_imp+1}": self.phantom_profiles["impurity_density"].sel(element=imp) - for num_imp, imp in enumerate(self.plasma.impurities)}} - result["PHANTOMS"] = {**result["PHANTOMS"], **{ + for num_imp, imp in enumerate(self.plasma.impurities) + }) + result["PHANTOMS"].update({ f"TIZ{num_imp + 1}": self.phantom_profiles["ion_temperature"].sel(element=imp) - for num_imp, imp in enumerate(self.plasma.impurities)}} - + for num_imp, imp in enumerate(self.plasma.impurities) + }) result["PROFILES"] = { "RHO_POLOIDAL": self.plasma.rho, diff --git a/indica/workflows/bayes_workflow_dev.py b/indica/workflows/bayes_workflow_dev.py index ba0e97ed..5fa34615 100644 --- a/indica/workflows/bayes_workflow_dev.py +++ b/indica/workflows/bayes_workflow_dev.py @@ -15,7 +15,7 @@ from indica.converters.line_of_sight import LineOfSightTransform from indica.workflows.abstract_bayes_workflow import AbstractBayesWorkflow -from indica.writers.bda_tree import create_nodes, write_nodes +from indica.writers.bda_tree import create_nodes, write_nodes, check_analysis_run # global configurations DEFAULT_PHANTOM_PARAMS = { @@ -107,6 +107,7 @@ def __init__( self, pulse=None, pulse_to_write=None, + run="RUN01", diagnostics=None, param_names=None, opt_quantity=None, @@ -130,6 +131,7 @@ def __init__( ): self.pulse = pulse self.pulse_to_write = pulse_to_write + self.run = run self.diagnostics = diagnostics self.param_names = param_names self.opt_quantity = opt_quantity @@ -369,6 +371,7 @@ def _exp_data(self): def __call__(self, filepath = "./results/test/", **kwargs): if self.mds_write: + check_analysis_run(self.pulse, self.run) self.node_structure = create_nodes(pulse_to_write=self.pulse_to_write, diagnostic_quantities=self.opt_quantity, mode="NEW") @@ -391,6 +394,7 @@ def __call__(self, filepath = "./results/test/", **kwargs): run = DevBayesWorkflow( pulse=10009, pulse_to_write=23000101, + run="RUN01", diagnostics=["xrcs", "efit", "smmh1", "cxff_pi"], opt_quantity=OPTIMISED_QUANTITY, param_names=OPTIMISED_PARAMS, From 9f62a78bde2e100f5a8e5f7a20acb5c9dcb7f2a5 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Fri, 14 Jul 2023 12:46:55 +0100 Subject: [PATCH 081/160] stashing --- indica/workflows/bayes_workflow_example.py | 24 +++++++++++++++++----- 1 file changed, 19 insertions(+), 5 deletions(-) diff --git a/indica/workflows/bayes_workflow_example.py b/indica/workflows/bayes_workflow_example.py index ea954787..4093c5c0 100644 --- a/indica/workflows/bayes_workflow_example.py +++ b/indica/workflows/bayes_workflow_example.py @@ -345,17 +345,31 @@ def __call__(self, filepath = "./results/test/", **kwargs): if __name__ == "__main__": run = ExampleWorkflow( pulse=10009, + pulse_to_write=None, dt=0.005, tsample=0.060, - diagnostics=["efit", "smmh1", "cxff_pi"], - param_names=OPTIMISED_PARAMS, - opt_quantity=OPTIMISED_QUANTITY, - priors=DEFAULT_PRIORS, + diagnostics=None, + param_names=None, + opt_quantity=None, + priors=None, + phantoms=False, + phantom_params=None, + model_kwargs=None, iterations=20, nwalkers=50, burn_frac=0.10, - sample_high_density=True + tstart=0.02, + tend=0.10, + dt=0.01, + tsample=0.06, + + sample_high_density=True, + + + mds_write=False, + plot=True, + sample_high_density=False, ) test = run(filepath="./results/test/") From 3828c49f4058bf09610c99304049af576d0c106c Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Fri, 14 Jul 2023 13:31:32 +0100 Subject: [PATCH 082/160] fixing key names --- indica/workflows/bayes_plots.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/indica/workflows/bayes_plots.py b/indica/workflows/bayes_plots.py index dc9ed77d..a4e06773 100644 --- a/indica/workflows/bayes_plots.py +++ b/indica/workflows/bayes_plots.py @@ -292,7 +292,7 @@ def plot_bayes_result( figheader=figheader, ylabel="Energy [J]", ) - key = "smmh1.ne" + key = "SMMH1.NE" if key in model_data.keys(): violinplot( model_data, @@ -362,7 +362,7 @@ def plot_bayes_result( ) key = "TI" plot_profile( - profiles[key].sel(element="ar"), + profiles[key], key, figheader=figheader, filename="temperature", @@ -384,7 +384,7 @@ def plot_bayes_result( ) key = "NI" plot_profile( - profiles[key].sel(element="h"), + profiles[key], key, figheader=figheader, filetype=filetype, From 310a440c9ffb36f2c55355edbf5f5687e3afa046 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Fri, 14 Jul 2023 13:31:58 +0100 Subject: [PATCH 083/160] updating and testing example --- indica/workflows/bayes_workflow_example.py | 124 ++++++++++++--------- 1 file changed, 69 insertions(+), 55 deletions(-) diff --git a/indica/workflows/bayes_workflow_example.py b/indica/workflows/bayes_workflow_example.py index 4093c5c0..da6cbb22 100644 --- a/indica/workflows/bayes_workflow_example.py +++ b/indica/workflows/bayes_workflow_example.py @@ -1,15 +1,12 @@ -import pickle - import emcee import numpy as np -import xarray as xr -import pandas as pd import flatdict +import copy from scipy.stats import loguniform -from pathlib import Path from indica.readers.read_st40 import ReadST40 from indica.bayesmodels import BayesModels, get_uniform +from indica.workflows.bayes_plots import plot_bayes_result from indica.models.interferometry import Interferometry from indica.models.helike_spectroscopy import Helike_spectroscopy from indica.models.charge_exchange import ChargeExchange @@ -17,11 +14,10 @@ from indica.models.plasma import Plasma from indica.converters.line_of_sight import LineOfSightTransform -from abstract_bayes_workflow import AbstractBayesWorkflow +from indica.workflows.abstract_bayes_workflow import AbstractBayesWorkflow +from indica.writers.bda_tree import create_nodes, write_nodes, check_analysis_run # global configurations -from indica.workflows.bayes_plots import plot_bayes_result - DEFAULT_PHANTOM_PARAMS = { "Ne_prof.y0": 5e19, "Ne_prof.wcenter": 0.4, @@ -54,7 +50,7 @@ "Ne_prof.wcenter": get_uniform(0.1, 0.8), "Ne_prof.peaking": get_uniform(1, 6), "ar_conc": loguniform(0.0001, 0.01), - "Nimp_prof.y0": get_uniform(1e16, 1e18), + "Nimp_prof.y0": loguniform(1e16, 1e18), "Nimp_prof.y1": get_uniform(1e15, 2e16), "Ne_prof.y0/Nimp_prof.y0": lambda x1, x2: np.where( (x1 > x2 * 100) & (x1 < x2 * 1e5), 1, 0 @@ -99,17 +95,26 @@ # "Ti_prof.peaking", ] -OPTIMISED_QUANTITY = ["cxff_pi.ti", "efit.wp", "smmh1.ne"] +OPTIMISED_QUANTITY = [ + "xrcs.spectra", + "cxff_pi.ti", + "efit.wp", + "smmh1.ne"] -class ExampleWorkflow(AbstractBayesWorkflow): +class DevBayesWorkflow(AbstractBayesWorkflow): def __init__( self, pulse=None, + pulse_to_write=None, + run="RUN01", + diagnostics=None, param_names=None, opt_quantity=None, priors = None, - diagnostics=None, + phantoms=False, + phantom_params=None, + model_kwargs = None, nwalkers=50, tstart=0.02, @@ -119,37 +124,43 @@ def __init__( iterations=100, burn_frac=0, - phantoms=False, + mds_write=False, + plot=True, sample_high_density = False, ): - self.pulse = pulse + self.pulse_to_write = pulse_to_write + self.run = run + self.diagnostics = diagnostics self.param_names = param_names self.opt_quantity = opt_quantity self.priors = priors - self.diagnostics = diagnostics + self.phantom_params = phantom_params + self.model_kwargs = model_kwargs + self.phantoms = phantoms self.tstart = tstart self.tend = tend self.dt = dt self.tsample = tsample - self.nwalkers = nwalkers self.iterations = iterations self.burn_frac = burn_frac - self.phantoms = phantoms + + self.mds_write = mds_write + self.plot = plot self.sample_high_density = sample_high_density - for attribute in ["pulse", "param_names", "opt_quantity", "priors", "diagnostics"]: + for attribute in ["pulse", "param_names", "opt_quantity", "priors", "diagnostics", "phantom_params"]: if getattr(self, attribute) is None: raise ValueError(f"{attribute} needs to be defined") self.setup_plasma() self.save_phantom_profiles() - self.read_data(diagnostics) + self.read_data(self.diagnostics) self.setup_opt_data(phantoms=self.phantoms) - self.setup_models(diagnostics) - self.setup_optimiser() + self.setup_models(self.diagnostics) + self.setup_optimiser(self.model_kwargs) def setup_plasma(self): self.plasma = Plasma( @@ -168,7 +179,7 @@ def setup_plasma(self): self.plasma.time_to_calculate = self.plasma.t[ np.abs(self.tsample - self.plasma.t).argmin() ] - self.plasma.update_profiles(DEFAULT_PHANTOM_PARAMS) + self.plasma.update_profiles(self.phantom_params) self.plasma.build_atomic_data(calc_power_loss=False) def setup_opt_data(self, phantoms=False): @@ -200,7 +211,6 @@ def setup_models(self, diagnostics: list): los_transform.set_equilibrium(self.plasma.equilibrium) model = Interferometry(name=diag) model.set_los_transform(los_transform) - model.plasma = self.plasma elif diag == "xrcs": los_transform = self.data["xrcs"]["te_kw"].transform @@ -213,24 +223,23 @@ def setup_models(self, diagnostics: list): , ) model.set_los_transform(los_transform) - model.plasma = self.plasma elif diag == "efit": model = EquilibriumReconstruction(name="efit") - model.plasma = self.plasma elif diag == "cxff_pi": transform = self.data[diag]["ti"].transform transform.set_equilibrium(self.plasma.equilibrium) model = ChargeExchange(name=diag, element="ar") model.set_transect_transform(transform) - model.plasma = self.plasma else: raise ValueError(f"{diag} not found in setup_models") + model.plasma = self.plasma + self.models[diag] = model - def setup_optimiser(self, ): + def setup_optimiser(self, model_kwargs): self.bayesopt = BayesModels( plasma=self.plasma, @@ -248,6 +257,7 @@ def setup_optimiser(self, ): log_prob_fn=self.bayesopt.ln_posterior, parameter_names=self.param_names, moves=self.move, + kwargs=model_kwargs, ) if self.sample_high_density: @@ -335,45 +345,49 @@ def _exp_data(self): ) def __call__(self, filepath = "./results/test/", **kwargs): + + if self.mds_write: + # check_analysis_run(self.pulse, self.run) + self.node_structure = create_nodes(pulse_to_write=self.pulse_to_write, + diagnostic_quantities=self.opt_quantity, + mode="NEW") + self.run_sampler() self.save_pickle(filepath=filepath) - plot_bayes_result(self.result, filepath) - return self.result + if self.plot: # currently requires result + plot_bayes_result(self.result, filepath) + + self.result = self.dict_of_dataarray_to_numpy(self.result) + if self.mds_write: + write_nodes(self.pulse_to_write, self.node_structure, self.result) + + return self.result if __name__ == "__main__": - run = ExampleWorkflow( + + run = DevBayesWorkflow( pulse=10009, - pulse_to_write=None, + pulse_to_write=23000101, + run="RUN01", + diagnostics=["xrcs", "efit", "smmh1", "cxff_pi"], + opt_quantity=OPTIMISED_QUANTITY, + param_names=OPTIMISED_PARAMS, + phantom_params=DEFAULT_PHANTOM_PARAMS, + priors=DEFAULT_PRIORS, + + iterations=5, + nwalkers=20, + burn_frac=0.10, dt=0.005, tsample=0.060, - diagnostics=None, - param_names=None, - opt_quantity=None, - priors=None, - phantoms=False, - phantom_params=None, - model_kwargs=None, - - iterations=20, - nwalkers=50, - burn_frac=0.10, - tstart=0.02, - tend=0.10, - dt=0.01, - tsample=0.06, - - sample_high_density=True, - - mds_write=False, + mds_write=True, plot=True, + phantoms=False, sample_high_density=False, + model_kwargs= { "background": 100} ) + results = run(filepath="./results/test/",) - test = run(filepath="./results/test/") - - flattest = flatdict.FlatDict(test, delimiter=".") - for key in flattest.keys(): - print(key) From d8f9be559593a4515139e091936ef47cba5f6362 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Fri, 14 Jul 2023 13:42:55 +0100 Subject: [PATCH 084/160] added background as attribute of helike model --- indica/models/helike_spectroscopy.py | 7 ++++--- indica/workflows/bayes_workflow_example.py | 9 ++++++--- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/indica/models/helike_spectroscopy.py b/indica/models/helike_spectroscopy.py index 3079ddd5..98c34813 100644 --- a/indica/models/helike_spectroscopy.py +++ b/indica/models/helike_spectroscopy.py @@ -48,6 +48,7 @@ def __init__( window_lim: list = [0.394, 0.401], window_vector=None, window_masks: list = [], + background = None, ): """ Read all atomic data and initialise objects @@ -87,6 +88,7 @@ def __init__( else: self.pecs = self._set_adas_pecs() + self.background = background self.window_masks = window_masks self.window_vector = window_vector if self.window_vector is not None: @@ -581,7 +583,6 @@ def __call__( calc_rho: bool = False, minimum_lines: bool = False, moment_analysis: bool = False, - background=None, **kwargs, ): """ @@ -692,8 +693,8 @@ def __call__( ) if moment_analysis: self._calculate_temperatures() - if background is not None: - self.measured_spectra = self.measured_spectra + background + if self.background is not None: + self.measured_spectra = self.measured_spectra + self.background self._build_bckc_dictionary() return self.bckc diff --git a/indica/workflows/bayes_workflow_example.py b/indica/workflows/bayes_workflow_example.py index da6cbb22..756c5fcc 100644 --- a/indica/workflows/bayes_workflow_example.py +++ b/indica/workflows/bayes_workflow_example.py @@ -213,14 +213,17 @@ def setup_models(self, diagnostics: list): model.set_los_transform(los_transform) elif diag == "xrcs": + + self.xrcs_background = self.opt_data["xrcs.spectra"].sel(wavelength=slice(0.392, 0.388)).mean( + dim="wavelength").sel(t=self.plasma.time_to_calculate) los_transform = self.data["xrcs"]["te_kw"].transform model = Helike_spectroscopy( name="xrcs", window_masks=[slice(0.394, 0.396)], window_vector=self.data[diag][ "spectra" - ].wavelength.values - , + ].wavelength.values, + background=self.xrcs_background, ) model.set_los_transform(los_transform) @@ -387,7 +390,7 @@ def __call__(self, filepath = "./results/test/", **kwargs): plot=True, phantoms=False, sample_high_density=False, - model_kwargs= { "background": 100} + model_kwargs= {} ) results = run(filepath="./results/test/",) From 1f0a0f9f8f48a8acb916c25e09eea58aacd38be6 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Fri, 14 Jul 2023 15:43:37 +0100 Subject: [PATCH 085/160] fixed phantom methods --- indica/workflows/bayes_workflow_dev.py | 27 +++-- indica/workflows/bayes_workflow_example.py | 113 ++++++++++----------- 2 files changed, 69 insertions(+), 71 deletions(-) diff --git a/indica/workflows/bayes_workflow_dev.py b/indica/workflows/bayes_workflow_dev.py index 5fa34615..bf7521c1 100644 --- a/indica/workflows/bayes_workflow_dev.py +++ b/indica/workflows/bayes_workflow_dev.py @@ -74,25 +74,25 @@ } OPTIMISED_PARAMS = [ - # "Ne_prof.y1", + "Ne_prof.y1", "Ne_prof.y0", - # "Ne_prof.peaking", + "Ne_prof.peaking", # "Ne_prof.wcenter", # "Ne_prof.wped", # "ar_conc", # "Nimp_prof.y1", "Nimp_prof.y0", # "Nimp_prof.wcenter", - # "Nimp_prof.wped", - # "Nimp_prof.peaking", + "Nimp_prof.wped", + "Nimp_prof.peaking", "Te_prof.y0", # "Te_prof.wped", - # "Te_prof.wcenter", - # "Te_prof.peaking", + "Te_prof.wcenter", + "Te_prof.peaking", "Ti_prof.y0", # "Ti_prof.wped", - # "Ti_prof.wcenter", - # "Ti_prof.peaking", + "Ti_prof.wcenter", + "Ti_prof.peaking", ] OPTIMISED_QUANTITY = [ @@ -371,18 +371,17 @@ def _exp_data(self): def __call__(self, filepath = "./results/test/", **kwargs): if self.mds_write: - check_analysis_run(self.pulse, self.run) + # check_analysis_run(self.pulse, self.run) self.node_structure = create_nodes(pulse_to_write=self.pulse_to_write, diagnostic_quantities=self.opt_quantity, mode="NEW") self.run_sampler() self.save_pickle(filepath=filepath) - self.result = self.dict_of_dataarray_to_numpy(self.result) - if self.plot: plot_bayes_result(self.result, filepath) + self.result = self.dict_of_dataarray_to_numpy(self.result) if self.mds_write: write_nodes(self.pulse_to_write, self.node_structure, self.result) @@ -401,14 +400,14 @@ def __call__(self, filepath = "./results/test/", **kwargs): phantom_params=DEFAULT_PHANTOM_PARAMS, priors=DEFAULT_PRIORS, - iterations=5, - nwalkers=20, + iterations=100, + nwalkers=50, burn_frac=0.10, dt=0.005, tsample=0.060, mds_write=True, - plot=False, + plot=True, phantoms=False, sample_high_density=False, model_kwargs= { "background": 100} diff --git a/indica/workflows/bayes_workflow_example.py b/indica/workflows/bayes_workflow_example.py index 756c5fcc..e476504c 100644 --- a/indica/workflows/bayes_workflow_example.py +++ b/indica/workflows/bayes_workflow_example.py @@ -74,25 +74,25 @@ } OPTIMISED_PARAMS = [ - # "Ne_prof.y1", + "Ne_prof.y1", "Ne_prof.y0", - # "Ne_prof.peaking", + "Ne_prof.peaking", # "Ne_prof.wcenter", # "Ne_prof.wped", # "ar_conc", # "Nimp_prof.y1", "Nimp_prof.y0", # "Nimp_prof.wcenter", - # "Nimp_prof.wped", - # "Nimp_prof.peaking", + "Nimp_prof.wped", + "Nimp_prof.peaking", "Te_prof.y0", # "Te_prof.wped", - # "Te_prof.wcenter", - # "Te_prof.peaking", + "Te_prof.wcenter", + "Te_prof.peaking", "Ti_prof.y0", # "Ti_prof.wped", - # "Ti_prof.wcenter", - # "Ti_prof.peaking", + "Ti_prof.wcenter", + "Ti_prof.peaking", ] OPTIMISED_QUANTITY = [ @@ -102,7 +102,7 @@ "smmh1.ne"] -class DevBayesWorkflow(AbstractBayesWorkflow): +class ExampleBayesWorkflow(AbstractBayesWorkflow): def __init__( self, pulse=None, @@ -158,8 +158,8 @@ def __init__( self.setup_plasma() self.save_phantom_profiles() self.read_data(self.diagnostics) - self.setup_opt_data(phantoms=self.phantoms) self.setup_models(self.diagnostics) + self.setup_opt_data(phantoms=self.phantoms) self.setup_optimiser(self.model_kwargs) def setup_plasma(self): @@ -171,7 +171,7 @@ def setup_plasma(self): impurities=("ar", "c"), impurity_concentration=( 0.001, - 0.02, + 0.04, ), full_run=False, n_rad=20, @@ -182,12 +182,6 @@ def setup_plasma(self): self.plasma.update_profiles(self.phantom_params) self.plasma.build_atomic_data(calc_power_loss=False) - def setup_opt_data(self, phantoms=False): - if phantoms: - self._phantom_data() - else: - self._exp_data() - def setup_models(self, diagnostics: list): self.models = {} for diag in self.diagnostics: @@ -213,17 +207,13 @@ def setup_models(self, diagnostics: list): model.set_los_transform(los_transform) elif diag == "xrcs": - - self.xrcs_background = self.opt_data["xrcs.spectra"].sel(wavelength=slice(0.392, 0.388)).mean( - dim="wavelength").sel(t=self.plasma.time_to_calculate) los_transform = self.data["xrcs"]["te_kw"].transform model = Helike_spectroscopy( name="xrcs", window_masks=[slice(0.394, 0.396)], window_vector=self.data[diag][ "spectra" - ].wavelength.values, - background=self.xrcs_background, + ].wavelength.values * 0.1, ) model.set_los_transform(los_transform) @@ -242,33 +232,11 @@ def setup_models(self, diagnostics: list): self.models[diag] = model - def setup_optimiser(self, model_kwargs): - - self.bayesopt = BayesModels( - plasma=self.plasma, - data=self.opt_data, - diagnostic_models=[*self.models.values()], - quant_to_optimise=self.opt_quantity, - priors=self.priors, - ) - - ndim = len(self.param_names) - self.move = [(emcee.moves.StretchMove(), 0.9), (emcee.moves.DEMove(), 0.1)] - self.sampler = emcee.EnsembleSampler( - self.nwalkers, - ndim, - log_prob_fn=self.bayesopt.ln_posterior, - parameter_names=self.param_names, - moves=self.move, - kwargs=model_kwargs, - ) - - if self.sample_high_density: - self.start_points = self.bayesopt.sample_from_high_density_region(self.param_names, self.sampler, self.nwalkers) + def setup_opt_data(self, phantoms=False): + if phantoms: + self._phantom_data() else: - self.start_points = self.bayesopt.sample_from_priors( - self.param_names, size=self.nwalkers - ) + self._exp_data() def _phantom_data(self, noise=False, noise_factor=0.1): self.opt_data = {} @@ -279,19 +247,22 @@ def _phantom_data(self, noise=False, noise_factor=0.1): .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) ) if "xrcs" in self.diagnostics: + self.opt_data["xrcs.spectra"] = ( self.models["xrcs"]() .pop("spectra") .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) ) - self.opt_data["xrcs.background"] = None + self.opt_data["xrcs.spectra"]["error"] = np.sqrt( + self.opt_data["xrcs.spectra"] + ) if "cxff_pi" in self.diagnostics: cxrs_data = ( self.models["cxff_pi"]() .pop("ti") .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) ) - self.opt_data["cxff_pi.ti"] = cxrs_data.where(cxrs_data.channel == 2) + self.opt_data["cxff_pi.ti"] = cxrs_data.where(cxrs_data.channel == 2, drop=True) if "efit" in self.diagnostics: self.opt_data["efit.wp"] = ( self.models["efit"]() @@ -337,7 +308,7 @@ def _exp_data(self): & (self.opt_data["xrcs.spectra"].wavelength > 0.388), drop=True, ) - self.opt_data["xrcs.background"] = background.mean(dim="wavelength") + self.models["xrcs"].background = background.mean(dim="wavelength").sel(t=self.plasma.time_to_calculate) self.opt_data["xrcs.spectra"]["error"] = np.sqrt( self.opt_data["xrcs.spectra"] + background.std(dim="wavelength") ** 2 ) @@ -347,6 +318,34 @@ def _exp_data(self): self.opt_data["cxff_pi"]["ti"].channel == 2, drop=True ) + def setup_optimiser(self, model_kwargs): + + self.bayesopt = BayesModels( + plasma=self.plasma, + data=self.opt_data, + diagnostic_models=[*self.models.values()], + quant_to_optimise=self.opt_quantity, + priors=self.priors, + ) + + ndim = len(self.param_names) + self.move = [(emcee.moves.StretchMove(), 0.9), (emcee.moves.DEMove(), 0.1)] + self.sampler = emcee.EnsembleSampler( + self.nwalkers, + ndim, + log_prob_fn=self.bayesopt.ln_posterior, + parameter_names=self.param_names, + moves=self.move, + kwargs=model_kwargs, + ) + + if self.sample_high_density: + self.start_points = self.bayesopt.sample_from_high_density_region(self.param_names, self.sampler, self.nwalkers) + else: + self.start_points = self.bayesopt.sample_from_priors( + self.param_names, size=self.nwalkers + ) + def __call__(self, filepath = "./results/test/", **kwargs): if self.mds_write: @@ -358,7 +357,7 @@ def __call__(self, filepath = "./results/test/", **kwargs): self.run_sampler() self.save_pickle(filepath=filepath) - if self.plot: # currently requires result + if self.plot: # currently requires result with DataArrays plot_bayes_result(self.result, filepath) self.result = self.dict_of_dataarray_to_numpy(self.result) @@ -370,7 +369,7 @@ def __call__(self, filepath = "./results/test/", **kwargs): if __name__ == "__main__": - run = DevBayesWorkflow( + run = ExampleBayesWorkflow( pulse=10009, pulse_to_write=23000101, run="RUN01", @@ -380,16 +379,16 @@ def __call__(self, filepath = "./results/test/", **kwargs): phantom_params=DEFAULT_PHANTOM_PARAMS, priors=DEFAULT_PRIORS, - iterations=5, - nwalkers=20, + iterations=200, + nwalkers=50, burn_frac=0.10, dt=0.005, tsample=0.060, mds_write=True, plot=True, - phantoms=False, - sample_high_density=False, + phantoms=True, + sample_high_density=True, model_kwargs= {} ) results = run(filepath="./results/test/",) From 7623a5449d5d6b543fc614593addc091c676fa2c Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Fri, 14 Jul 2023 15:44:02 +0100 Subject: [PATCH 086/160] rearranged methods --- indica/workflows/abstract_bayes_workflow.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/indica/workflows/abstract_bayes_workflow.py b/indica/workflows/abstract_bayes_workflow.py index 86d1fb99..b4282ccd 100644 --- a/indica/workflows/abstract_bayes_workflow.py +++ b/indica/workflows/abstract_bayes_workflow.py @@ -27,8 +27,8 @@ def __init__(self, self.setup_plasma() self.save_phantom_profiles() self.read_data(self.diagnostics) - self.setup_opt_data(self.phantoms) self.setup_models(self.diagnostics) + self.setup_opt_data(self.phantoms) self.setup_optimiser() @abstractmethod @@ -47,14 +47,6 @@ def read_data(self, diagnostics: list): self.plasma.set_equilibrium(self.reader.equilibrium) self.data = self.reader.binned_data - - @abstractmethod - def setup_opt_data(self, phantom: bool = False): - """ - Prepare the data in necessary format for optimiser i.e. flat dictionary - """ - self.opt_data = {} - @abstractmethod def setup_models(self, diagnostics: list): """ @@ -63,6 +55,13 @@ def setup_models(self, diagnostics: list): """ self.models = {} + @abstractmethod + def setup_opt_data(self, phantom: bool = False): + """ + Prepare the data in necessary format for optimiser i.e. flat dictionary + """ + self.opt_data = {} + @abstractmethod def setup_optimiser(self, model_kwargs): """ From b3b7a264f53fb2e8c4ed28a8dcf461e6b52dc250 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Wed, 2 Aug 2023 10:24:52 +0100 Subject: [PATCH 087/160] moved background to call kwarg --- indica/models/helike_spectroscopy.py | 7 +++---- tests/unit/workflows/__init__.py | 0 2 files changed, 3 insertions(+), 4 deletions(-) create mode 100644 tests/unit/workflows/__init__.py diff --git a/indica/models/helike_spectroscopy.py b/indica/models/helike_spectroscopy.py index 98c34813..c4e1b869 100644 --- a/indica/models/helike_spectroscopy.py +++ b/indica/models/helike_spectroscopy.py @@ -48,7 +48,6 @@ def __init__( window_lim: list = [0.394, 0.401], window_vector=None, window_masks: list = [], - background = None, ): """ Read all atomic data and initialise objects @@ -88,7 +87,6 @@ def __init__( else: self.pecs = self._set_adas_pecs() - self.background = background self.window_masks = window_masks self.window_vector = window_vector if self.window_vector is not None: @@ -583,6 +581,7 @@ def __call__( calc_rho: bool = False, minimum_lines: bool = False, moment_analysis: bool = False, + background: int = None, **kwargs, ): """ @@ -693,8 +692,8 @@ def __call__( ) if moment_analysis: self._calculate_temperatures() - if self.background is not None: - self.measured_spectra = self.measured_spectra + self.background + if background is not None: + self.measured_spectra = self.measured_spectra + background self._build_bckc_dictionary() return self.bckc diff --git a/tests/unit/workflows/__init__.py b/tests/unit/workflows/__init__.py new file mode 100644 index 00000000..e69de29b From 4a7e022981e11dbd924d2ec94376259ca733f181 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Wed, 2 Aug 2023 11:20:52 +0100 Subject: [PATCH 088/160] params and kwargs now are given as model_{var} and model prefix is removed before model call --- indica/bayesmodels.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/indica/bayesmodels.py b/indica/bayesmodels.py index d6092f09..019e18e3 100644 --- a/indica/bayesmodels.py +++ b/indica/bayesmodels.py @@ -61,7 +61,7 @@ def __init__( if missing_data: # gives list of keys in quant_to_optimise but not data raise ValueError(f"{missing_data} not found in data given") - def _build_bckc(self, params: dict, **kwargs): + def _build_bckc(self, params, **kwargs): """ TODO: consider how to handle when models have overlapping kwargs Parameters @@ -73,13 +73,19 @@ def _build_bckc(self, params: dict, **kwargs): ------- bckc of results """ - # Params is a dictionary which is updated by optimiser, - # kwargs is constant i.e. settings for models self.bckc: dict = {} for model in self.diagnostic_models: - self.bckc = dict( - self.bckc, **{model.name: {**model(**{**params, **kwargs})}} - ) + # removes "model.name_" from params and kwargs then passes them to model e.g. xrcs_background -> background + _nuisance_params = {param_name.replace(model.name+"_", ""): + param_value for param_name, param_value in params.items() + if model.name in param_name} + _model_kwargs = {kwarg_name.replace(model.name+"_", ""): + kwarg_value for kwarg_name, kwarg_value in kwargs.items() + if model.name in kwarg_name} + + model_bckc = {model.name: {**model(**{**_nuisance_params, **_model_kwargs})}} + self.bckc = dict(self.bckc, **model_bckc) + self.bckc = flatdict.FlatDict(self.bckc, delimiter=".") return From e78f63814a34e8b7c58fca08bcc015e585055206 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Wed, 2 Aug 2023 13:19:09 +0100 Subject: [PATCH 089/160] adding _phantom_data and _exp_data abstract methods --- indica/workflows/abstract_bayes_workflow.py | 37 ++++++++++++++------- 1 file changed, 25 insertions(+), 12 deletions(-) diff --git a/indica/workflows/abstract_bayes_workflow.py b/indica/workflows/abstract_bayes_workflow.py index b4282ccd..a8754f6d 100644 --- a/indica/workflows/abstract_bayes_workflow.py +++ b/indica/workflows/abstract_bayes_workflow.py @@ -24,29 +24,29 @@ def __init__(self, self.opt_quantity = opt_quantity self.priors = priors + self.read_data(self.diagnostics) self.setup_plasma() self.save_phantom_profiles() - self.read_data(self.diagnostics) self.setup_models(self.diagnostics) self.setup_opt_data(self.phantoms) self.setup_optimiser() - @abstractmethod - def setup_plasma(self): - """ - Contains all methods and settings for plasma object to be used in optimisation - """ - self.plasma = None - def read_data(self, diagnostics: list): self.reader = ReadST40( self.pulse, tstart=self.tstart, tend=self.tend, dt=self.dt ) self.reader(diagnostics) - self.plasma.set_equilibrium(self.reader.equilibrium) self.data = self.reader.binned_data + @abstractmethod + def setup_plasma(self): + """ + Contains all methods and settings for setting up / initialising plasma object + """ + self.plasma = None + self.plasma.set_equilibrium(self.reader.equilibrium) + @abstractmethod def setup_models(self, diagnostics: list): """ @@ -55,12 +55,25 @@ def setup_models(self, diagnostics: list): """ self.models = {} + @abstractmethod + def _phantom_data(self): + opt_data = {} + return opt_data + + @abstractmethod + def _exp_data(self): + opt_data = {} + return opt_data + @abstractmethod def setup_opt_data(self, phantom: bool = False): """ - Prepare the data in necessary format for optimiser i.e. flat dictionary + Get and prepare the data in necessary format for optimiser """ - self.opt_data = {} + if phantom: + self.opt_data = self._phantom_data() + else: + self.opt_data = self._exp_data() @abstractmethod def setup_optimiser(self, model_kwargs): @@ -82,7 +95,7 @@ def save_phantom_profiles(self, kinetic_profiles=None): t=self.plasma.time_to_calculate) * 0 for profile_key in kinetic_profiles} - self.phantom_profiles = phantom_profiles + self.phantom_profiles = phantom_profiles def _build_result_dict(self): From 82d6fee53c41c21ea51fe8c64a44aa05f22cbaea Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Wed, 2 Aug 2023 13:20:28 +0100 Subject: [PATCH 090/160] renaming to BayesWorkflowExample --- indica/workflows/bayes_workflow_example.py | 169 ++++++++++----------- 1 file changed, 81 insertions(+), 88 deletions(-) diff --git a/indica/workflows/bayes_workflow_example.py b/indica/workflows/bayes_workflow_example.py index e476504c..6b598060 100644 --- a/indica/workflows/bayes_workflow_example.py +++ b/indica/workflows/bayes_workflow_example.py @@ -74,59 +74,60 @@ } OPTIMISED_PARAMS = [ - "Ne_prof.y1", + # "Ne_prof.y1", "Ne_prof.y0", - "Ne_prof.peaking", + # "Ne_prof.peaking", # "Ne_prof.wcenter", # "Ne_prof.wped", # "ar_conc", # "Nimp_prof.y1", "Nimp_prof.y0", # "Nimp_prof.wcenter", - "Nimp_prof.wped", - "Nimp_prof.peaking", + # "Nimp_prof.wped", + # "Nimp_prof.peaking", "Te_prof.y0", # "Te_prof.wped", "Te_prof.wcenter", "Te_prof.peaking", "Ti_prof.y0", - # "Ti_prof.wped", + "Ti_prof.wped", "Ti_prof.wcenter", "Ti_prof.peaking", ] OPTIMISED_QUANTITY = [ - "xrcs.spectra", - "cxff_pi.ti", - "efit.wp", - "smmh1.ne"] + "xrcs.spectra", + "cxff_pi.ti", + "efit.wp", + "smmh1.ne", +] -class ExampleBayesWorkflow(AbstractBayesWorkflow): +class BayesWorkflowExample(AbstractBayesWorkflow): def __init__( self, - pulse=None, - pulse_to_write=None, - run="RUN01", - diagnostics=None, - param_names=None, - opt_quantity=None, - priors = None, - phantoms=False, - phantom_params=None, - model_kwargs = None, - - nwalkers=50, - tstart=0.02, - tend=0.10, - dt=0.01, - tsample=0.06, - iterations=100, - burn_frac=0, - - mds_write=False, - plot=True, - sample_high_density = False, + pulse: int = None, + pulse_to_write: bool = None, + run: str = "RUN01", + diagnostics: list = None, + param_names: list = None, + opt_quantity: list = None, + priors: dict = None, + profile_params: dict = None, + phantoms: bool = False, + model_kwargs: dict = {}, + + nwalkers: int = 50, + tstart: float = 0.02, + tend: float = 0.10, + dt: float = 0.01, + tsample: float = 0.06, + iterations: int = 100, + burn_frac: float = 0, + + mds_write: bool = False, + plot: bool = True, + sample_high_density: bool = False, ): self.pulse = pulse self.pulse_to_write = pulse_to_write @@ -135,7 +136,7 @@ def __init__( self.param_names = param_names self.opt_quantity = opt_quantity self.priors = priors - self.phantom_params = phantom_params + self.profile_params = profile_params self.model_kwargs = model_kwargs self.phantoms = phantoms @@ -151,13 +152,13 @@ def __init__( self.plot = plot self.sample_high_density = sample_high_density - for attribute in ["pulse", "param_names", "opt_quantity", "priors", "diagnostics", "phantom_params"]: + for attribute in ["pulse", "param_names", "opt_quantity", "priors", "diagnostics", "profile_params"]: if getattr(self, attribute) is None: raise ValueError(f"{attribute} needs to be defined") + self.read_data(self.diagnostics) self.setup_plasma() self.save_phantom_profiles() - self.read_data(self.diagnostics) self.setup_models(self.diagnostics) self.setup_opt_data(phantoms=self.phantoms) self.setup_optimiser(self.model_kwargs) @@ -179,7 +180,8 @@ def setup_plasma(self): self.plasma.time_to_calculate = self.plasma.t[ np.abs(self.tsample - self.plasma.t).argmin() ] - self.plasma.update_profiles(self.phantom_params) + self.plasma.set_equilibrium(self.reader.equilibrium) + self.plasma.update_profiles(self.profile_params) self.plasma.build_atomic_data(calc_power_loss=False) def setup_models(self, diagnostics: list): @@ -228,33 +230,30 @@ def setup_models(self, diagnostics: list): else: raise ValueError(f"{diag} not found in setup_models") model.plasma = self.plasma - - self.models[diag] = model def setup_opt_data(self, phantoms=False): if phantoms: - self._phantom_data() + self.opt_data = self._phantom_data() else: - self._exp_data() + self.opt_data = self._exp_data() def _phantom_data(self, noise=False, noise_factor=0.1): - self.opt_data = {} + opt_data = {} if "smmh1" in self.diagnostics: - self.opt_data["smmh1.ne"] = ( + opt_data["smmh1.ne"] = ( self.models["smmh1"]() .pop("ne") .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) ) if "xrcs" in self.diagnostics: - - self.opt_data["xrcs.spectra"] = ( + opt_data["xrcs.spectra"] = ( self.models["xrcs"]() .pop("spectra") .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) ) - self.opt_data["xrcs.spectra"]["error"] = np.sqrt( - self.opt_data["xrcs.spectra"] + opt_data["xrcs.spectra"]["error"] = np.sqrt( + opt_data["xrcs.spectra"] ) if "cxff_pi" in self.diagnostics: cxrs_data = ( @@ -262,61 +261,56 @@ def _phantom_data(self, noise=False, noise_factor=0.1): .pop("ti") .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) ) - self.opt_data["cxff_pi.ti"] = cxrs_data.where(cxrs_data.channel == 2, drop=True) + opt_data["cxff_pi.ti"] = cxrs_data.where(cxrs_data.channel == 2, drop=True) if "efit" in self.diagnostics: - self.opt_data["efit.wp"] = ( + opt_data["efit.wp"] = ( self.models["efit"]() .pop("wp") .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) ) if noise: - self.opt_data["smmh1.ne"] = self.opt_data["smmh1.ne"] + self.opt_data[ + opt_data["smmh1.ne"] = opt_data["smmh1.ne"] + self.opt_data[ "smmh1.ne" ].max().values * np.random.normal(0, noise_factor, None) - self.opt_data["xrcs.spectra"] = self.opt_data[ - "xrcs.spectra" - ] + np.random.normal( + opt_data["xrcs.spectra"] = opt_data[ + "xrcs.spectra" + ] + np.random.normal( 0, - np.sqrt(self.opt_data["xrcs.spectra"].values[0,]), - self.opt_data["xrcs.spectra"].shape[1], + np.sqrt(opt_data["xrcs.spectra"].values[0,]), + opt_data["xrcs.spectra"].shape[1], ) - self.opt_data["cxff_pi.ti"] = self.opt_data[ - "cxff_pi.ti" - ] + self.opt_data["cxff_pi.ti"].max().values * np.random.normal( - 0, noise_factor, self.opt_data["cxff_pi.ti"].shape[1] + opt_data["cxff_pi.ti"] = opt_data[ + "cxff_pi.ti" + ] + opt_data["cxff_pi.ti"].max().values * np.random.normal( + 0, noise_factor, opt_data["cxff_pi.ti"].shape[1] ) - self.opt_data["efit.wp"] = self.opt_data["efit.wp"] + self.opt_data[ + opt_data["efit.wp"] = opt_data["efit.wp"] + opt_data[ "efit.wp" ].max().values * np.random.normal(0, noise_factor, None) - - self.phantom_profiles = {} - for key in ["electron_density", "impurity_density", "electron_temperature", "ion_temperature", - "ion_density", "fast_density", "neutral_density"]: - self.phantom_profiles[key] = getattr(self.plasma, key).sel( - t=self.plasma.time_to_calculate - ).copy() + return opt_data def _exp_data(self): - self.opt_data = flatdict.FlatDict(self.data, ".") + opt_data = flatdict.FlatDict(self.data, ".") if "xrcs" in self.diagnostics: - self.opt_data["xrcs.spectra"]["wavelength"] = ( - self.opt_data["xrcs.spectra"].wavelength * 0.1 + opt_data["xrcs.spectra"]["wavelength"] = ( + opt_data["xrcs.spectra"].wavelength * 0.1 ) - background = self.opt_data["xrcs.spectra"].where( - (self.opt_data["xrcs.spectra"].wavelength < 0.392) - & (self.opt_data["xrcs.spectra"].wavelength > 0.388), + background = opt_data["xrcs.spectra"].where( + (opt_data["xrcs.spectra"].wavelength < 0.392) + & (opt_data["xrcs.spectra"].wavelength > 0.388), drop=True, ) - self.models["xrcs"].background = background.mean(dim="wavelength").sel(t=self.plasma.time_to_calculate) - self.opt_data["xrcs.spectra"]["error"] = np.sqrt( - self.opt_data["xrcs.spectra"] + background.std(dim="wavelength") ** 2 + self.model_kwargs["xrcs_background"] = background.mean(dim="wavelength").sel(t=self.plasma.time_to_calculate) + opt_data["xrcs.spectra"]["error"] = np.sqrt( + opt_data["xrcs.spectra"] + background.std(dim="wavelength") ** 2 ) if "cxff_pi" in self.diagnostics: - self.opt_data["cxff_pi"]["ti"] = self.opt_data["cxff_pi"]["ti"].where( - self.opt_data["cxff_pi"]["ti"].channel == 2, drop=True + opt_data["cxff_pi"]["ti"] = opt_data["cxff_pi"]["ti"].where( + opt_data["cxff_pi"]["ti"].channel == 2, drop=True ) + return opt_data def setup_optimiser(self, model_kwargs): @@ -340,13 +334,14 @@ def setup_optimiser(self, model_kwargs): ) if self.sample_high_density: - self.start_points = self.bayesopt.sample_from_high_density_region(self.param_names, self.sampler, self.nwalkers) + self.start_points = self.bayesopt.sample_from_high_density_region(self.param_names, self.sampler, + self.nwalkers) else: self.start_points = self.bayesopt.sample_from_priors( self.param_names, size=self.nwalkers ) - def __call__(self, filepath = "./results/test/", **kwargs): + def __call__(self, filepath="./results/test/", **kwargs): if self.mds_write: # check_analysis_run(self.pulse, self.run) @@ -368,19 +363,19 @@ def __call__(self, filepath = "./results/test/", **kwargs): if __name__ == "__main__": - - run = ExampleBayesWorkflow( + run = BayesWorkflowExample( pulse=10009, pulse_to_write=23000101, run="RUN01", diagnostics=["xrcs", "efit", "smmh1", "cxff_pi"], opt_quantity=OPTIMISED_QUANTITY, param_names=OPTIMISED_PARAMS, - phantom_params=DEFAULT_PHANTOM_PARAMS, + profile_params=DEFAULT_PHANTOM_PARAMS, priors=DEFAULT_PRIORS, + model_kwargs={"xrcs_moment_analysis":False, }, - iterations=200, - nwalkers=50, + iterations=10, + nwalkers=20, burn_frac=0.10, dt=0.005, tsample=0.060, @@ -388,8 +383,6 @@ def __call__(self, filepath = "./results/test/", **kwargs): mds_write=True, plot=True, phantoms=True, - sample_high_density=True, - model_kwargs= {} + sample_high_density=False, ) - results = run(filepath="./results/test/",) - + results = run(filepath="./results/test/", ) From 2c1c48aa3157a90771d44028a8afa08b05b363f7 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Wed, 2 Aug 2023 13:22:43 +0100 Subject: [PATCH 091/160] initial commit --- .../workflows/test_bayes_workflow_example.py | 61 +++++++++++++++++++ 1 file changed, 61 insertions(+) create mode 100644 tests/unit/workflows/test_bayes_workflow_example.py diff --git a/tests/unit/workflows/test_bayes_workflow_example.py b/tests/unit/workflows/test_bayes_workflow_example.py new file mode 100644 index 00000000..f08ca42e --- /dev/null +++ b/tests/unit/workflows/test_bayes_workflow_example.py @@ -0,0 +1,61 @@ +from indica.workflows.bayes_workflow_example import ExampleBayesWorkflow + + +""" +IDEAS: + +including abstract class tests in here... + +if save phantom profiles is mutable + +does it initialise + +setup models: + does it run + give wrong diagnostic name + give duplicate names + give correct names + +_phantom_data: + does it run +check dictionary: + all have time dim + all exist + + with noise + all have error? + shape + +_exp_data: + does it run + test if no data + test shape of data + all have dim t + +setup_optimiser + does it run + start points + contains data + dims of start points + does sample from high density work + +call + does it run + are results all np.array + +""" + +class TestBayesWorkflowExample: + def setup_class(self): + return + + def setup_method(self): + return + + def teardown_method(self): + return + + def test_plasma_initializes(self): + assert hasattr(self, "plasma") + + From c7a2ca58b87ec257b5a536a1415b1b220b176ba5 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Wed, 2 Aug 2023 13:24:22 +0100 Subject: [PATCH 092/160] renamed phantom_params to profile_params --- indica/workflows/bayes_workflow_example.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/indica/workflows/bayes_workflow_example.py b/indica/workflows/bayes_workflow_example.py index 6b598060..b7733de3 100644 --- a/indica/workflows/bayes_workflow_example.py +++ b/indica/workflows/bayes_workflow_example.py @@ -18,7 +18,7 @@ from indica.writers.bda_tree import create_nodes, write_nodes, check_analysis_run # global configurations -DEFAULT_PHANTOM_PARAMS = { +DEFAULT_PROFILE_PARAMS = { "Ne_prof.y0": 5e19, "Ne_prof.wcenter": 0.4, "Ne_prof.peaking": 2, @@ -370,7 +370,7 @@ def __call__(self, filepath="./results/test/", **kwargs): diagnostics=["xrcs", "efit", "smmh1", "cxff_pi"], opt_quantity=OPTIMISED_QUANTITY, param_names=OPTIMISED_PARAMS, - profile_params=DEFAULT_PHANTOM_PARAMS, + profile_params=DEFAULT_PROFILE_PARAMS, priors=DEFAULT_PRIORS, model_kwargs={"xrcs_moment_analysis":False, }, From db25e708890df372c71ee308321a9ac4e385119f Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Wed, 2 Aug 2023 13:56:05 +0100 Subject: [PATCH 093/160] renamed bayesopt -> bayesmodel --- indica/workflows/bayes_workflow_example.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/indica/workflows/bayes_workflow_example.py b/indica/workflows/bayes_workflow_example.py index b7733de3..7739723f 100644 --- a/indica/workflows/bayes_workflow_example.py +++ b/indica/workflows/bayes_workflow_example.py @@ -314,7 +314,7 @@ def _exp_data(self): def setup_optimiser(self, model_kwargs): - self.bayesopt = BayesModels( + self.bayesmodel = BayesModels( plasma=self.plasma, data=self.opt_data, diagnostic_models=[*self.models.values()], @@ -327,7 +327,7 @@ def setup_optimiser(self, model_kwargs): self.sampler = emcee.EnsembleSampler( self.nwalkers, ndim, - log_prob_fn=self.bayesopt.ln_posterior, + log_prob_fn=self.bayesmodel.ln_posterior, parameter_names=self.param_names, moves=self.move, kwargs=model_kwargs, @@ -373,6 +373,7 @@ def __call__(self, filepath="./results/test/", **kwargs): profile_params=DEFAULT_PROFILE_PARAMS, priors=DEFAULT_PRIORS, model_kwargs={"xrcs_moment_analysis":False, }, + phantoms=True, iterations=10, nwalkers=20, @@ -382,7 +383,6 @@ def __call__(self, filepath="./results/test/", **kwargs): mds_write=True, plot=True, - phantoms=True, sample_high_density=False, ) results = run(filepath="./results/test/", ) From b1c5bd375a749e32dff8795d26d913326368c488 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Wed, 2 Aug 2023 13:57:24 +0100 Subject: [PATCH 094/160] renamed bayesopt -> bayesmodel --- indica/workflows/bayes_workflow_example.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/indica/workflows/bayes_workflow_example.py b/indica/workflows/bayes_workflow_example.py index 7739723f..170a08f3 100644 --- a/indica/workflows/bayes_workflow_example.py +++ b/indica/workflows/bayes_workflow_example.py @@ -375,8 +375,8 @@ def __call__(self, filepath="./results/test/", **kwargs): model_kwargs={"xrcs_moment_analysis":False, }, phantoms=True, - iterations=10, - nwalkers=20, + iterations=200, + nwalkers=50, burn_frac=0.10, dt=0.005, tsample=0.060, From f78ce649941226ca21377f2ad00c3f1d6e2fa700 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Wed, 2 Aug 2023 15:00:19 +0100 Subject: [PATCH 095/160] moved start point sampling to its own method _sample_start_points --- indica/workflows/bayes_workflow_example.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/indica/workflows/bayes_workflow_example.py b/indica/workflows/bayes_workflow_example.py index 170a08f3..cfe72309 100644 --- a/indica/workflows/bayes_workflow_example.py +++ b/indica/workflows/bayes_workflow_example.py @@ -332,8 +332,10 @@ def setup_optimiser(self, model_kwargs): moves=self.move, kwargs=model_kwargs, ) + self.start_points = self._sample_start_points(sample_high_density=self.sample_high_density) - if self.sample_high_density: + def _sample_start_points(self, sample_high_density: bool=True): + if sample_high_density: self.start_points = self.bayesopt.sample_from_high_density_region(self.param_names, self.sampler, self.nwalkers) else: From c3a418bd90cfcfeb66ae86786a348ef985fd5fa0 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Wed, 2 Aug 2023 15:04:35 +0100 Subject: [PATCH 096/160] fixed variable name --- indica/workflows/bayes_workflow_example.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/indica/workflows/bayes_workflow_example.py b/indica/workflows/bayes_workflow_example.py index cfe72309..98b68e8a 100644 --- a/indica/workflows/bayes_workflow_example.py +++ b/indica/workflows/bayes_workflow_example.py @@ -336,10 +336,10 @@ def setup_optimiser(self, model_kwargs): def _sample_start_points(self, sample_high_density: bool=True): if sample_high_density: - self.start_points = self.bayesopt.sample_from_high_density_region(self.param_names, self.sampler, + self.start_points = self.bayesmodel.sample_from_high_density_region(self.param_names, self.sampler, self.nwalkers) else: - self.start_points = self.bayesopt.sample_from_priors( + self.start_points = self.bayesmodel.sample_from_priors( self.param_names, size=self.nwalkers ) From 196eb60f7c978f6cf896dda6ea1899090c137332 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Wed, 2 Aug 2023 15:09:49 +0100 Subject: [PATCH 097/160] fixed bug with start_points being overwritten --- indica/workflows/bayes_workflow_example.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/indica/workflows/bayes_workflow_example.py b/indica/workflows/bayes_workflow_example.py index 98b68e8a..9f610183 100644 --- a/indica/workflows/bayes_workflow_example.py +++ b/indica/workflows/bayes_workflow_example.py @@ -336,12 +336,13 @@ def setup_optimiser(self, model_kwargs): def _sample_start_points(self, sample_high_density: bool=True): if sample_high_density: - self.start_points = self.bayesmodel.sample_from_high_density_region(self.param_names, self.sampler, + start_points = self.bayesmodel.sample_from_high_density_region(self.param_names, self.sampler, self.nwalkers) else: - self.start_points = self.bayesmodel.sample_from_priors( + start_points = self.bayesmodel.sample_from_priors( self.param_names, size=self.nwalkers ) + return start_points def __call__(self, filepath="./results/test/", **kwargs): From 31542b0835d80af6f0c3c47e0c11768ddcce84e1 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 8 Aug 2023 09:44:21 +0100 Subject: [PATCH 098/160] removed workflow_dev --- indica/workflows/bayes_workflow_dev.py | 416 ------------------------- 1 file changed, 416 deletions(-) delete mode 100644 indica/workflows/bayes_workflow_dev.py diff --git a/indica/workflows/bayes_workflow_dev.py b/indica/workflows/bayes_workflow_dev.py deleted file mode 100644 index bf7521c1..00000000 --- a/indica/workflows/bayes_workflow_dev.py +++ /dev/null @@ -1,416 +0,0 @@ -import emcee -import numpy as np -import flatdict -import copy -from scipy.stats import loguniform - -from indica.readers.read_st40 import ReadST40 -from indica.bayesmodels import BayesModels, get_uniform -from indica.workflows.bayes_plots import plot_bayes_result -from indica.models.interferometry import Interferometry -from indica.models.helike_spectroscopy import Helike_spectroscopy -from indica.models.charge_exchange import ChargeExchange -from indica.models.equilibrium_reconstruction import EquilibriumReconstruction -from indica.models.plasma import Plasma -from indica.converters.line_of_sight import LineOfSightTransform - -from indica.workflows.abstract_bayes_workflow import AbstractBayesWorkflow -from indica.writers.bda_tree import create_nodes, write_nodes, check_analysis_run - -# global configurations -DEFAULT_PHANTOM_PARAMS = { - "Ne_prof.y0": 5e19, - "Ne_prof.wcenter": 0.4, - "Ne_prof.peaking": 2, - "Ne_prof.y1": 2e18, - "Ne_prof.yend": 1e18, - "Ne_prof.wped": 2, - "Nimp_prof.y0": 3e16, - "Nimp_prof.y1": 0.5e16, - "Nimp_prof.wcenter": 0.4, - "Nimp_prof.wped": 6, - "Nimp_prof.peaking": 2, - "Te_prof.y0": 3000, - "Te_prof.y1": 50, - "Te_prof.wcenter": 0.4, - "Te_prof.wped": 3, - "Te_prof.peaking": 2, - "Ti_prof.y0": 6000, - "Ti_prof.y1": 50, - "Ti_prof.wcenter": 0.4, - "Ti_prof.wped": 3, - "Ti_prof.peaking": 2, -} - -DEFAULT_PRIORS = { - "Ne_prof.y0": get_uniform(2e19, 4e20), - "Ne_prof.y1": get_uniform(1e18, 1e19), - "Ne_prof.y0/Ne_prof.y1": lambda x1, x2: np.where((x1 > x2 * 2), 1, 0), - "Ne_prof.wped": get_uniform(1, 6), - "Ne_prof.wcenter": get_uniform(0.1, 0.8), - "Ne_prof.peaking": get_uniform(1, 6), - "ar_conc": loguniform(0.0001, 0.01), - "Nimp_prof.y0": loguniform(1e16, 1e18), - "Nimp_prof.y1": get_uniform(1e15, 2e16), - "Ne_prof.y0/Nimp_prof.y0": lambda x1, x2: np.where( - (x1 > x2 * 100) & (x1 < x2 * 1e5), 1, 0 - ), - "Nimp_prof.y0/Nimp_prof.y1": lambda x1, x2: np.where((x1 > x2), 1, 0), - "Nimp_prof.wped": get_uniform(1, 6), - "Nimp_prof.wcenter": get_uniform(0.1, 0.8), - "Nimp_prof.peaking": get_uniform(1, 20), - "Nimp_prof.peaking/Ne_prof.peaking": lambda x1, x2: np.where( - (x1 > x2), 1, 0 - ), # impurity always more peaked - "Te_prof.y0": get_uniform(1000, 5000), - "Te_prof.wped": get_uniform(1, 6), - "Te_prof.wcenter": get_uniform(0.1, 0.6), - "Te_prof.peaking": get_uniform(1, 6), - "Ti_prof.y0/Te_prof.y0": lambda x1, x2: np.where(x1 > x2, 1, 0), # hot ion mode - "Ti_prof.y0": get_uniform(3000, 10000), - "Ti_prof.wped": get_uniform(1, 6), - "Ti_prof.wcenter": get_uniform(0.1, 0.6), - "Ti_prof.peaking": get_uniform(1, 20), -} - -OPTIMISED_PARAMS = [ - "Ne_prof.y1", - "Ne_prof.y0", - "Ne_prof.peaking", - # "Ne_prof.wcenter", - # "Ne_prof.wped", - # "ar_conc", - # "Nimp_prof.y1", - "Nimp_prof.y0", - # "Nimp_prof.wcenter", - "Nimp_prof.wped", - "Nimp_prof.peaking", - "Te_prof.y0", - # "Te_prof.wped", - "Te_prof.wcenter", - "Te_prof.peaking", - "Ti_prof.y0", - # "Ti_prof.wped", - "Ti_prof.wcenter", - "Ti_prof.peaking", -] - -OPTIMISED_QUANTITY = [ - "xrcs.spectra", - "cxff_pi.ti", - "efit.wp", - "smmh1.ne"] - - -class DevBayesWorkflow(AbstractBayesWorkflow): - def __init__( - self, - pulse=None, - pulse_to_write=None, - run="RUN01", - diagnostics=None, - param_names=None, - opt_quantity=None, - priors = None, - phantoms=False, - phantom_params=None, - model_kwargs = None, - - nwalkers=50, - tstart=0.02, - tend=0.10, - dt=0.01, - tsample=0.06, - iterations=100, - burn_frac=0, - - mds_write=False, - plot=True, - sample_high_density = False, - fast_particles = False, - ): - self.pulse = pulse - self.pulse_to_write = pulse_to_write - self.run = run - self.diagnostics = diagnostics - self.param_names = param_names - self.opt_quantity = opt_quantity - self.priors = priors - self.phantom_params = phantom_params - self.model_kwargs = model_kwargs - self.phantoms = phantoms - - self.tstart = tstart - self.tend = tend - self.dt = dt - self.tsample = tsample - self.nwalkers = nwalkers - self.iterations = iterations - self.burn_frac = burn_frac - - self.mds_write = mds_write - self.plot = plot - self.sample_high_density = sample_high_density - self.fast_particles = fast_particles - - for attribute in ["pulse", "param_names", "opt_quantity", "priors", "diagnostics", "phantom_params"]: - if getattr(self, attribute) is None: - raise ValueError(f"{attribute} needs to be defined") - - self.setup_plasma() - self.save_phantom_profiles() - self.read_data(self.diagnostics) - self.setup_opt_data(phantoms=self.phantoms) - self.setup_models(self.diagnostics) - self.setup_optimiser(self.model_kwargs) - - def setup_plasma(self): - self.plasma = Plasma( - tstart=self.tstart, - tend=self.tend, - dt=self.dt, - main_ion="h", - impurities=("ar", "c"), - impurity_concentration=( - 0.001, - 0.02, - ), - full_run=False, - n_rad=20, - ) - self.plasma.time_to_calculate = self.plasma.t[ - np.abs(self.tsample - self.plasma.t).argmin() - ] - self.plasma.update_profiles(self.phantom_params) - self.plasma.build_atomic_data(calc_power_loss=False) - if self.fast_particles: - self._init_fast_particles() - - def _init_fast_particles(self): - st40_code = ReadST40(13110009, self.tstart, self.tend, dt=self.dt, tree="astra") - st40_code.get_raw_data("", "astra", "RUN573") - st40_code.bin_data_in_time(["astra"], self.tstart, self.tend, self.dt) - code_data = st40_code.binned_data["astra"] - Nf = ( - code_data["nf"].interp(rho_poloidal=self.plasma.rho, t=self.plasma.t) - * 1.0e19 - ) - self.plasma.fast_density.values = Nf.values - Nn = ( - code_data["nn"].interp(rho_poloidal=self.plasma.rho, t=self.plasma.t) - * 1.0e19 - ) - self.plasma.neutral_density.values = Nn.values - Pblon = code_data["pblon"].interp(rho_poloidal=self.plasma.rho, t=self.plasma.t) - self.plasma.pressure_fast_parallel.values = Pblon.values - Pbper = code_data["pbper"].interp(rho_poloidal=self.plasma.rho, t=self.plasma.t) - self.plasma.pressure_fast_perpendicular.values = Pbper.values - - def setup_opt_data(self, phantoms=False): - if phantoms: - self._phantom_data() - else: - self._exp_data() - - def setup_models(self, diagnostics: list): - self.models = {} - for diag in self.diagnostics: - if diag == "smmh1": - # los_transform = self.data["smmh1"]["ne"].transform - machine_dims = self.plasma.machine_dimensions - origin = np.array([[-0.38063365, 0.91893092, 0.01]]) - # end = np.array([[0, 0, 0.01]]) - direction = np.array([[0.38173721, -0.92387953, -0.02689453]]) - los_transform = LineOfSightTransform( - origin[:, 0], - origin[:, 1], - origin[:, 2], - direction[:, 0], - direction[:, 1], - direction[:, 2], - name="", - machine_dimensions=machine_dims, - passes=2, - ) - los_transform.set_equilibrium(self.plasma.equilibrium) - model = Interferometry(name=diag) - model.set_los_transform(los_transform) - - elif diag == "xrcs": - los_transform = self.data["xrcs"]["te_kw"].transform - model = Helike_spectroscopy( - name="xrcs", - window_masks=[slice(0.394, 0.396)], - window_vector=self.data[diag][ - "spectra" - ].wavelength.values - , - ) - model.set_los_transform(los_transform) - - elif diag == "efit": - model = EquilibriumReconstruction(name="efit") - - elif diag == "cxff_pi": - transform = self.data[diag]["ti"].transform - transform.set_equilibrium(self.plasma.equilibrium) - model = ChargeExchange(name=diag, element="ar") - model.set_transect_transform(transform) - else: - raise ValueError(f"{diag} not found in setup_models") - model.plasma = self.plasma - - - self.models[diag] = model - - def setup_optimiser(self, model_kwargs): - - self.bayesopt = BayesModels( - plasma=self.plasma, - data=self.opt_data, - diagnostic_models=[*self.models.values()], - quant_to_optimise=self.opt_quantity, - priors=self.priors, - ) - - ndim = len(self.param_names) - self.move = [(emcee.moves.StretchMove(), 0.9), (emcee.moves.DEMove(), 0.1)] - self.sampler = emcee.EnsembleSampler( - self.nwalkers, - ndim, - log_prob_fn=self.bayesopt.ln_posterior, - parameter_names=self.param_names, - moves=self.move, - kwargs=model_kwargs, - ) - - if self.sample_high_density: - self.start_points = self.bayesopt.sample_from_high_density_region(self.param_names, self.sampler, self.nwalkers) - else: - self.start_points = self.bayesopt.sample_from_priors( - self.param_names, size=self.nwalkers - ) - - def _phantom_data(self, noise=False, noise_factor=0.1): - self.opt_data = {} - if "smmh1" in self.diagnostics: - self.opt_data["smmh1.ne"] = ( - self.models["smmh1"]() - .pop("ne") - .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) - ) - if "xrcs" in self.diagnostics: - self.opt_data["xrcs.spectra"] = ( - self.models["xrcs"]() - .pop("spectra") - .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) - ) - self.opt_data["xrcs.background"] = None - if "cxff_pi" in self.diagnostics: - cxrs_data = ( - self.models["cxff_pi"]() - .pop("ti") - .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) - ) - self.opt_data["cxff_pi.ti"] = cxrs_data.where(cxrs_data.channel == 2) - if "efit" in self.diagnostics: - self.opt_data["efit.wp"] = ( - self.models["efit"]() - .pop("wp") - .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) - ) - - if noise: - self.opt_data["smmh1.ne"] = self.opt_data["smmh1.ne"] + self.opt_data[ - "smmh1.ne" - ].max().values * np.random.normal(0, noise_factor, None) - self.opt_data["xrcs.spectra"] = self.opt_data[ - "xrcs.spectra" - ] + np.random.normal( - 0, - np.sqrt(self.opt_data["xrcs.spectra"].values[0,]), - self.opt_data["xrcs.spectra"].shape[1], - ) - self.opt_data["cxff_pi.ti"] = self.opt_data[ - "cxff_pi.ti" - ] + self.opt_data["cxff_pi.ti"].max().values * np.random.normal( - 0, noise_factor, self.opt_data["cxff_pi.ti"].shape[1] - ) - self.opt_data["efit.wp"] = self.opt_data["efit.wp"] + self.opt_data[ - "efit.wp" - ].max().values * np.random.normal(0, noise_factor, None) - - self.phantom_profiles = {} - for key in ["electron_density", "impurity_density", "electron_temperature", "ion_temperature", - "ion_density", "fast_density", "neutral_density"]: - self.phantom_profiles[key] = getattr(self.plasma, key).sel( - t=self.plasma.time_to_calculate - ).copy() - - def _exp_data(self): - self.opt_data = flatdict.FlatDict(self.data, ".") - if "xrcs" in self.diagnostics: - self.opt_data["xrcs.spectra"]["wavelength"] = ( - self.opt_data["xrcs.spectra"].wavelength * 0.1 - ) - background = self.opt_data["xrcs.spectra"].where( - (self.opt_data["xrcs.spectra"].wavelength < 0.392) - & (self.opt_data["xrcs.spectra"].wavelength > 0.388), - drop=True, - ) - self.opt_data["xrcs.background"] = background.mean(dim="wavelength") - self.opt_data["xrcs.spectra"]["error"] = np.sqrt( - self.opt_data["xrcs.spectra"] + background.std(dim="wavelength") ** 2 - ) - - if "cxff_pi" in self.diagnostics: - self.opt_data["cxff_pi"]["ti"] = self.opt_data["cxff_pi"]["ti"].where( - self.opt_data["cxff_pi"]["ti"].channel == 2, drop=True - ) - - def __call__(self, filepath = "./results/test/", **kwargs): - - if self.mds_write: - # check_analysis_run(self.pulse, self.run) - self.node_structure = create_nodes(pulse_to_write=self.pulse_to_write, - diagnostic_quantities=self.opt_quantity, - mode="NEW") - - self.run_sampler() - self.save_pickle(filepath=filepath) - if self.plot: - plot_bayes_result(self.result, filepath) - - self.result = self.dict_of_dataarray_to_numpy(self.result) - if self.mds_write: - write_nodes(self.pulse_to_write, self.node_structure, self.result) - - return self.result - - -if __name__ == "__main__": - - run = DevBayesWorkflow( - pulse=10009, - pulse_to_write=23000101, - run="RUN01", - diagnostics=["xrcs", "efit", "smmh1", "cxff_pi"], - opt_quantity=OPTIMISED_QUANTITY, - param_names=OPTIMISED_PARAMS, - phantom_params=DEFAULT_PHANTOM_PARAMS, - priors=DEFAULT_PRIORS, - - iterations=100, - nwalkers=50, - burn_frac=0.10, - dt=0.005, - tsample=0.060, - - mds_write=True, - plot=True, - phantoms=False, - sample_high_density=False, - model_kwargs= { "background": 100} - ) - results = run(filepath="./results/test/",) - From cc218b407daa1da2807cdbc06353d15685f90a00 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 8 Aug 2023 09:51:28 +0100 Subject: [PATCH 099/160] moving equilibrium to workflow object --- indica/workflows/abstract_bayes_workflow.py | 1 + 1 file changed, 1 insertion(+) diff --git a/indica/workflows/abstract_bayes_workflow.py b/indica/workflows/abstract_bayes_workflow.py index a8754f6d..fcbc763e 100644 --- a/indica/workflows/abstract_bayes_workflow.py +++ b/indica/workflows/abstract_bayes_workflow.py @@ -37,6 +37,7 @@ def read_data(self, diagnostics: list): self.pulse, tstart=self.tstart, tend=self.tend, dt=self.dt ) self.reader(diagnostics) + self.equilibrium = self.reader.equilibrium self.data = self.reader.binned_data @abstractmethod From 9cb0266276da6d75f784dbbe18e2bf73bdd6ef86 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 8 Aug 2023 10:05:32 +0100 Subject: [PATCH 100/160] when reading raw data save transforms to their own attribute --- indica/readers/read_st40.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/indica/readers/read_st40.py b/indica/readers/read_st40.py index 61549965..07e95aec 100644 --- a/indica/readers/read_st40.py +++ b/indica/readers/read_st40.py @@ -75,6 +75,7 @@ def __init__( self.equilibrium: Equilibrium self.raw_data: dict = {} self.binned_data: dict = {} + self.transforms: dict = {} def reset_data(self): self.raw_data = {} @@ -104,6 +105,7 @@ def get_raw_data(self, uid: str, instrument: str, revision: RevisionLike = 0): transform = data[quant].transform if hasattr(transform, "set_equilibrium"): transform.set_equilibrium(self.equilibrium) + self.transforms[instrument] = transform self.raw_data[instrument] = data return data From 68c18722c354afe478c5bd0a4e70360257969e9b Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 8 Aug 2023 10:09:14 +0100 Subject: [PATCH 101/160] transforms saved to workflow class --- indica/workflows/abstract_bayes_workflow.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/indica/workflows/abstract_bayes_workflow.py b/indica/workflows/abstract_bayes_workflow.py index fcbc763e..25053151 100644 --- a/indica/workflows/abstract_bayes_workflow.py +++ b/indica/workflows/abstract_bayes_workflow.py @@ -31,13 +31,13 @@ def __init__(self, self.setup_opt_data(self.phantoms) self.setup_optimiser() - def read_data(self, diagnostics: list): self.reader = ReadST40( self.pulse, tstart=self.tstart, tend=self.tend, dt=self.dt ) self.reader(diagnostics) self.equilibrium = self.reader.equilibrium + self.transforms = self.reader.transforms self.data = self.reader.binned_data @abstractmethod From f0f561fa8271a66caf5eee602731529bfee5843e Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 8 Aug 2023 10:19:04 +0100 Subject: [PATCH 102/160] fixed name --- indica/models/helike_spectroscopy.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/indica/models/helike_spectroscopy.py b/indica/models/helike_spectroscopy.py index c4e1b869..6e334c72 100644 --- a/indica/models/helike_spectroscopy.py +++ b/indica/models/helike_spectroscopy.py @@ -28,7 +28,7 @@ } -class Helike_spectroscopy(DiagnosticModel): +class HelikeSpectrometer(DiagnosticModel): """ Data and methods to model XRCS spectrometer measurements From 3b2eaeee6f21aa5835adabc055daf68a17a2f3e2 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 8 Aug 2023 11:30:30 +0100 Subject: [PATCH 103/160] renamed example transform --- indica/models/helike_spectroscopy.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/indica/models/helike_spectroscopy.py b/indica/models/helike_spectroscopy.py index 6e334c72..d0a32d82 100644 --- a/indica/models/helike_spectroscopy.py +++ b/indica/models/helike_spectroscopy.py @@ -752,7 +752,7 @@ def select_transition(adf15_data, transition: str, wavelength: float): return pec -def helike_LOS_example(nchannels=3): +def helike_transform_example(nchannels=3): los_end = np.full((nchannels, 3), 0.0) los_end[:, 0] = 0.17 los_end[:, 1] = 0.0 @@ -783,9 +783,9 @@ def example_run(pulse: int = 9229, plasma=None, plot=False, **kwargs): ) # Create new diagnostic diagnostic_name = "xrcs" - los_transform = helike_LOS_example(3) + los_transform = helike_transform_example(3) los_transform.set_equilibrium(plasma.equilibrium) - model = Helike_spectroscopy( + model = HelikeSpectrometer( diagnostic_name, window_masks=[], ) From ae200d36fb69d7cfe28e6c0316d059a5d6a7e8b7 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 8 Aug 2023 11:34:42 +0100 Subject: [PATCH 104/160] made example los function --- indica/models/interferometry.py | 27 +++++++++++++++------------ 1 file changed, 15 insertions(+), 12 deletions(-) diff --git a/indica/models/interferometry.py b/indica/models/interferometry.py index f7c7458f..29a6b42b 100644 --- a/indica/models/interferometry.py +++ b/indica/models/interferometry.py @@ -88,20 +88,11 @@ def __call__( return self.bckc - -def example_run(pulse: int = None, plasma=None, plot=False): - if plasma is None: - plasma = example_plasma(pulse=pulse) - - # Create new interferometers diagnostics - diagnostic_name = "smmh1" +def smmh1_transform_example(): los_start = np.array([[0.8, 0, 0], [0.8, 0, -0.1], [0.8, 0, -0.2]]) los_end = np.array([[0.17, 0, 0], [0.17, 0, -0.25], [0.17, 0, -0.2]]) origin = los_start direction = los_end - los_start - model = Interferometry( - diagnostic_name, - ) los_transform = LineOfSightTransform( origin[:, 0], origin[:, 1], @@ -109,10 +100,22 @@ def example_run(pulse: int = None, plasma=None, plot=False): direction[:, 0], direction[:, 1], direction[:, 2], - name=diagnostic_name, - machine_dimensions=plasma.machine_dimensions, + name="smmh1", + machine_dimensions=((0.15, 0.95), (-0.7, 0.7)), passes=2, ) + return los_transform + +def example_run(pulse: int = None, plasma=None, plot=False): + if plasma is None: + plasma = example_plasma(pulse=pulse) + + # Create new interferometers diagnostics + diagnostic_name = "smmh1" + model = Interferometry( + diagnostic_name, + ) + los_transform = smmh1_transform_example() los_transform.set_equilibrium(plasma.equilibrium) model.set_los_transform(los_transform) model.set_plasma(plasma) From 3d9d54078bd78bc5404532d892c4f7a6d52ea4b0 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 8 Aug 2023 11:35:19 +0100 Subject: [PATCH 105/160] minor name fix --- indica/models/helike_spectroscopy.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/indica/models/helike_spectroscopy.py b/indica/models/helike_spectroscopy.py index d0a32d82..d4538b20 100644 --- a/indica/models/helike_spectroscopy.py +++ b/indica/models/helike_spectroscopy.py @@ -768,7 +768,7 @@ def helike_transform_example(nchannels=3): direction[0:nchannels, 0], direction[0:nchannels, 1], direction[0:nchannels, 2], - name="diagnostic_name", + name="xrcs", machine_dimensions=((0.15, 0.95), (-0.7, 0.7)), passes=1, ) From cd04fb1c82a447485c22af617f1d2d758594af31 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 8 Aug 2023 11:35:54 +0100 Subject: [PATCH 106/160] made read_test_data function --- indica/workflows/abstract_bayes_workflow.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/indica/workflows/abstract_bayes_workflow.py b/indica/workflows/abstract_bayes_workflow.py index 25053151..180a3443 100644 --- a/indica/workflows/abstract_bayes_workflow.py +++ b/indica/workflows/abstract_bayes_workflow.py @@ -6,7 +6,7 @@ import pickle from indica.readers.read_st40 import ReadST40 - +from indica.equilibrium import fake_equilibrium class AbstractBayesWorkflow(ABC): @abstractmethod @@ -31,6 +31,12 @@ def __init__(self, self.setup_opt_data(self.phantoms) self.setup_optimiser() + def read_test_data(self, diagnostic_transforms:dict): + # Used with phantom data for purposes of tests + self.equilibrium = fake_equilibrium(self.tstart, self.tend, self.dt, ) + self.transforms = diagnostic_transforms + self.data = {} + def read_data(self, diagnostics: list): self.reader = ReadST40( self.pulse, tstart=self.tstart, tend=self.tend, dt=self.dt From dfdd79107ad23f2315a780c1d1fb2f44c7906f0d Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 8 Aug 2023 11:39:53 +0100 Subject: [PATCH 107/160] made example los function --- indica/models/charge_exchange.py | 26 ++++++++++++++------------ 1 file changed, 14 insertions(+), 12 deletions(-) diff --git a/indica/models/charge_exchange.py b/indica/models/charge_exchange.py index a7cfadaf..1aae94ac 100644 --- a/indica/models/charge_exchange.py +++ b/indica/models/charge_exchange.py @@ -106,6 +106,19 @@ def __call__( return self.bckc +def pi_transform_example(nchannels:int): + x_positions = np.linspace(0.2, 0.8, nchannels) + y_positions = np.linspace(0.0, 0.0, nchannels) + z_positions = np.linspace(0.0, 0.0, nchannels) + + transect_transform = TransectCoordinates( + x_positions, + y_positions, + z_positions, + "pi", + machine_dimensions=((0.15, 0.95), (-0.7, 0.7)), + ) + return transect_transform def example_run( pulse: int = None, @@ -120,18 +133,7 @@ def example_run( plasma = example_plasma(pulse=pulse) # Create new interferometers diagnostics - nchannels = 5 - x_positions = np.linspace(0.2, 0.8, nchannels) - y_positions = np.linspace(0.0, 0.0, nchannels) - z_positions = np.linspace(0.0, 0.0, nchannels) - - transect_transform = TransectCoordinates( - x_positions, - y_positions, - z_positions, - diagnostic_name, - machine_dimensions=plasma.machine_dimensions, - ) + transect_transform = pi_transform_example(5) transect_transform.set_equilibrium(plasma.equilibrium) model = ChargeExchange( diagnostic_name, From 4c6c1c2ab783d3832e26289948500a2c5a1437ae Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 8 Aug 2023 11:51:09 +0100 Subject: [PATCH 108/160] minor type --- indica/workflows/abstract_bayes_workflow.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/indica/workflows/abstract_bayes_workflow.py b/indica/workflows/abstract_bayes_workflow.py index 180a3443..33f4ec2f 100644 --- a/indica/workflows/abstract_bayes_workflow.py +++ b/indica/workflows/abstract_bayes_workflow.py @@ -299,7 +299,7 @@ def run_sampler(self): for blob_name in blob_names } self.accept_frac = self.sampler.acceptance_fraction.sum() - self.prior_sample = self.bayesopt.sample_from_priors( + self.prior_sample = self.bayesmodel.sample_from_priors( self.param_names, size=int(1e4) ) self.post_sample = self.sampler.get_chain(flat=True) From 0ce85fabe13c4446d704df111a105dff0a3b4480 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 8 Aug 2023 11:54:06 +0100 Subject: [PATCH 109/160] added reader to read_test_data method --- indica/workflows/abstract_bayes_workflow.py | 1 + 1 file changed, 1 insertion(+) diff --git a/indica/workflows/abstract_bayes_workflow.py b/indica/workflows/abstract_bayes_workflow.py index 33f4ec2f..215742ff 100644 --- a/indica/workflows/abstract_bayes_workflow.py +++ b/indica/workflows/abstract_bayes_workflow.py @@ -33,6 +33,7 @@ def __init__(self, def read_test_data(self, diagnostic_transforms:dict): # Used with phantom data for purposes of tests + self.reader = None self.equilibrium = fake_equilibrium(self.tstart, self.tend, self.dt, ) self.transforms = diagnostic_transforms self.data = {} From fa5ac444ba42ef9d9984ad0dd0354bb8e905771e Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 8 Aug 2023 14:48:29 +0100 Subject: [PATCH 110/160] fixed copying workflow object states --- .../workflows/test_bayes_workflow_example.py | 151 ++++++++++++------ 1 file changed, 105 insertions(+), 46 deletions(-) diff --git a/tests/unit/workflows/test_bayes_workflow_example.py b/tests/unit/workflows/test_bayes_workflow_example.py index f08ca42e..a81c5cfe 100644 --- a/tests/unit/workflows/test_bayes_workflow_example.py +++ b/tests/unit/workflows/test_bayes_workflow_example.py @@ -1,61 +1,120 @@ -from indica.workflows.bayes_workflow_example import ExampleBayesWorkflow +import pytest +import copy +import numpy as np +from indica.workflows.bayes_workflow_example import BayesWorkflowExample, DEFAULT_PRIORS, DEFAULT_PROFILE_PARAMS, \ + OPTIMISED_PARAMS, OPTIMISED_QUANTITY """ -IDEAS: - -including abstract class tests in here... - -if save phantom profiles is mutable - -does it initialise - -setup models: - does it run - give wrong diagnostic name - give duplicate names - give correct names - -_phantom_data: - does it run -check dictionary: - all have time dim - all exist - - with noise - all have error? - shape - -_exp_data: - does it run - test if no data - test shape of data - all have dim t - -setup_optimiser - does it run - start points - contains data - dims of start points - does sample from high density work - -call - does it run - are results all np.array +TODO: +test kwarg handling + +Mock reader for testing experimental data methods """ + class TestBayesWorkflowExample: def setup_class(self): - return + + self.default_settings = dict(pulse=None, + pulse_to_write=23000111, + run="RUN01", + diagnostics=["xrcs", "efit", "smmh1", "cxff_pi"], + opt_quantity=OPTIMISED_QUANTITY, + param_names=OPTIMISED_PARAMS, + profile_params=DEFAULT_PROFILE_PARAMS, + priors=DEFAULT_PRIORS, + model_kwargs={"xrcs_moment_analysis": False, }, + phantoms=True, + + iterations=1, + nwalkers=20, + burn_frac=0.10, + dt=0.005, + tsample=0.060, + + mds_write=False, + plot=False, + sample_high_density=False, ) + + self.untouched_workflow = BayesWorkflowExample(**self.default_settings) + self.workflow = copy.deepcopy(self.untouched_workflow) def setup_method(self): - return + self.workflow = copy.deepcopy(self.untouched_workflow) def teardown_method(self): - return + self.workflow = None + + def test_workflow_initializes(self): + attributes_to_check = ["plasma", "phantom_profiles", "data", "reader", "opt_data", "models", "bayesmodel",] + for attribute in attributes_to_check: + if not hasattr(self.workflow, attribute): + raise ValueError(f"missing {attribute} in workflow object") + assert True + + def test_init_phantoms_false_with_example_plasma(self): + with pytest.raises(ValueError): + example = BayesWorkflowExample(dict(self.default_settings, **{"phantoms": False})) + + def test_init_not_including_all_required_inputs(self): + with pytest.raises(ValueError): + example = BayesWorkflowExample(dict(self.default_settings, **{"param_names": None})) + + # def test_reader_has_read_all_diagnostic_data(self): + # assert all(diag_name in self.workflow.reader.keys() for diag_name in self.workflow.diagnostics) + + def test_plasma_has_equilibrium(self): + assert hasattr(self.workflow.plasma, "equilibrium") + + def test_phantom_profiles_are_not_mutatable(self): + phantoms = copy.deepcopy(self.workflow.phantom_profiles) + self.workflow.plasma.electron_temperature += 1 + assert (phantoms is not self.workflow.phantom_profiles) + + def test_setup_models_with_wrong_diagnostic_names(self): + with pytest.raises(Exception): + self.workflow.setup_models(["foo", "bar", "xrcs"]) + + def test_phantom_data_exists(self): + self.workflow._phantom_data() + assert self.workflow.opt_data + + # def test_experimental_data_exists(self): + # self.workflow._exp_data() + # assert self.workflow.opt_data + + def test_phantom_data_has_time_dim(self): + self.workflow._phantom_data() + for key, value in self.workflow.opt_data.items(): + assert "t" in value.dims + + # def test_experimental_data_has_time_dim(self): + # self.workflow._exp_data() + # for key, value in self.workflow.opt_data.items(): + # assert "t" in value.dims + + def test_phantom_data_runs_with_noise_added(self): + self.workflow._phantom_data(noise=True) + assert self.workflow.opt_data + + def test_sampling_from_high_density(self): + self.workflow._sample_start_points(sample_high_density=True) + assert True + + def test_sampling_from_priors(self): + self.workflow._sample_start_points(sample_high_density=False) + assert True - def test_plasma_initializes(self): - assert hasattr(self, "plasma") + def test_worklow_has_results_after_run(self): + self.workflow.run_sampler() + if not hasattr(self.workflow, "result"): + raise ValueError(f"missing result in workflow object") + assert True +if __name__ == "__main__": + test = TestBayesWorkflowExample() + test.setup_class() + print() From 7d343a4cec6a41455fc11c1e94aa4b44004424c1 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 8 Aug 2023 14:57:00 +0100 Subject: [PATCH 111/160] updated example to work with assign_profiles --- indica/models/plasma.py | 15 ++------------- 1 file changed, 2 insertions(+), 13 deletions(-) diff --git a/indica/models/plasma.py b/indica/models/plasma.py index fcf79388..4ce7bac6 100644 --- a/indica/models/plasma.py +++ b/indica/models/plasma.py @@ -1257,6 +1257,7 @@ def example_run( impurity_concentration=impurity_concentration, full_run=full_run, verbose=verbose, + n_rad=n_rad, **kwargs, ) plasma.build_atomic_data(default=True, calc_power_loss=calc_power_loss) @@ -1272,27 +1273,15 @@ def example_run( nimp_wcenter = np.linspace(0.4, 0.1, nt) for i, t in enumerate(plasma.t): plasma.Te_prof.peaking = te_peaking[i] - plasma.assign_profiles(profile="electron_temperature", t=t) - plasma.Ti_prof.peaking = te_peaking[i] plasma.Ti_prof.y0 = ti0[i] - plasma.assign_profiles(profile="ion_temperature", t=t) - plasma.Vrot_prof.peaking = vrot_peaking[i] plasma.Vrot_prof.y0 = vrot0[i] - plasma.assign_profiles(profile="toroidal_rotation", t=t) - plasma.Ne_prof.peaking = ne_peaking[i] - plasma.assign_profiles(profile="electron_density", t=t) - plasma.Nimp_prof.peaking = nimp_peaking[i] plasma.Nimp_prof.y0 = nimp_y0[i] plasma.Nimp_prof.wcenter = nimp_wcenter[i] - for elem in plasma.impurities: - plasma.assign_profiles(profile="impurity_density", t=t, element=elem) - - for elem in plasma.elements: - plasma.assign_profiles(profile="toroidal_rotation", t=t, element=elem) + plasma.assign_profiles(impurity_to_profile="ar",) if pulse is None: equilibrium_data = fake_equilibrium_data( From 489d14a4735a465fb51da9728b1e9995cfa664c5 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 8 Aug 2023 14:57:49 +0100 Subject: [PATCH 112/160] can now run with pulse = None and synthetic transforms/equilibrium --- indica/workflows/bayes_workflow_example.py | 91 ++++++++++++---------- 1 file changed, 51 insertions(+), 40 deletions(-) diff --git a/indica/workflows/bayes_workflow_example.py b/indica/workflows/bayes_workflow_example.py index 9f610183..b2d5f4e3 100644 --- a/indica/workflows/bayes_workflow_example.py +++ b/indica/workflows/bayes_workflow_example.py @@ -1,22 +1,19 @@ import emcee import numpy as np import flatdict -import copy from scipy.stats import loguniform -from indica.readers.read_st40 import ReadST40 from indica.bayesmodels import BayesModels, get_uniform from indica.workflows.bayes_plots import plot_bayes_result -from indica.models.interferometry import Interferometry -from indica.models.helike_spectroscopy import Helike_spectroscopy -from indica.models.charge_exchange import ChargeExchange -from indica.models.equilibrium_reconstruction import EquilibriumReconstruction -from indica.models.plasma import Plasma -from indica.converters.line_of_sight import LineOfSightTransform - from indica.workflows.abstract_bayes_workflow import AbstractBayesWorkflow from indica.writers.bda_tree import create_nodes, write_nodes, check_analysis_run +from indica.models.interferometry import Interferometry, smmh1_transform_example +from indica.models.helike_spectroscopy import HelikeSpectrometer, helike_transform_example +from indica.models.charge_exchange import ChargeExchange, pi_transform_example +from indica.models.equilibrium_reconstruction import EquilibriumReconstruction +from indica.models.plasma import Plasma + # global configurations DEFAULT_PROFILE_PARAMS = { "Ne_prof.y0": 5e19, @@ -152,11 +149,20 @@ def __init__( self.plot = plot self.sample_high_density = sample_high_density - for attribute in ["pulse", "param_names", "opt_quantity", "priors", "diagnostics", "profile_params"]: + for attribute in ["param_names", "opt_quantity", "priors", "diagnostics", "profile_params"]: if getattr(self, attribute) is None: raise ValueError(f"{attribute} needs to be defined") - self.read_data(self.diagnostics) + if self.pulse is None and self.phantoms is False: + raise ValueError("Set phantoms to True when running phantom plasma i.e. pulse=None") + + if pulse is None: + print("Running in test mode") + self.read_test_data({"xrcs": helike_transform_example(1), + "smmh1": smmh1_transform_example(), + "cxff_pi": pi_transform_example(5)}) + else: + self.read_data(self.diagnostics) self.setup_plasma() self.save_phantom_profiles() self.setup_models(self.diagnostics) @@ -180,42 +186,46 @@ def setup_plasma(self): self.plasma.time_to_calculate = self.plasma.t[ np.abs(self.tsample - self.plasma.t).argmin() ] - self.plasma.set_equilibrium(self.reader.equilibrium) + self.plasma.set_equilibrium(self.equilibrium) self.plasma.update_profiles(self.profile_params) self.plasma.build_atomic_data(calc_power_loss=False) def setup_models(self, diagnostics: list): self.models = {} - for diag in self.diagnostics: + for diag in diagnostics: if diag == "smmh1": - # los_transform = self.data["smmh1"]["ne"].transform - machine_dims = self.plasma.machine_dimensions - origin = np.array([[-0.38063365, 0.91893092, 0.01]]) - # end = np.array([[0, 0, 0.01]]) - direction = np.array([[0.38173721, -0.92387953, -0.02689453]]) - los_transform = LineOfSightTransform( - origin[:, 0], - origin[:, 1], - origin[:, 2], - direction[:, 0], - direction[:, 1], - direction[:, 2], - name="", - machine_dimensions=machine_dims, - passes=2, - ) + los_transform = self.transforms[diag] + # machine_dims = self.plasma.machine_dimensions + # origin = np.array([[-0.38063365, 0.91893092, 0.01]]) + # # end = np.array([[0, 0, 0.01]]) + # direction = np.array([[0.38173721, -0.92387953, -0.02689453]]) + # los_transform = LineOfSightTransform( + # origin[:, 0], + # origin[:, 1], + # origin[:, 2], + # direction[:, 0], + # direction[:, 1], + # direction[:, 2], + # name="", + # machine_dimensions=machine_dims, + # passes=2, + # ) los_transform.set_equilibrium(self.plasma.equilibrium) model = Interferometry(name=diag) model.set_los_transform(los_transform) elif diag == "xrcs": - los_transform = self.data["xrcs"]["te_kw"].transform - model = Helike_spectroscopy( + los_transform = self.transforms[diag] + los_transform.set_equilibrium(self.plasma.equilibrium) + window_vector = None + if hasattr(self, "data"): + if diag in self.data.keys(): + window_vector = self.data[diag]["spectra"].wavelength.values * 0.1 + + model = HelikeSpectrometer( name="xrcs", window_masks=[slice(0.394, 0.396)], - window_vector=self.data[diag][ - "spectra" - ].wavelength.values * 0.1, + window_vector=window_vector ) model.set_los_transform(los_transform) @@ -223,7 +233,7 @@ def setup_models(self, diagnostics: list): model = EquilibriumReconstruction(name="efit") elif diag == "cxff_pi": - transform = self.data[diag]["ti"].transform + transform = self.transforms[diag] transform.set_equilibrium(self.plasma.equilibrium) model = ChargeExchange(name=diag, element="ar") model.set_transect_transform(transform) @@ -301,7 +311,8 @@ def _exp_data(self): & (opt_data["xrcs.spectra"].wavelength > 0.388), drop=True, ) - self.model_kwargs["xrcs_background"] = background.mean(dim="wavelength").sel(t=self.plasma.time_to_calculate) + self.model_kwargs["xrcs_background"] = background.mean(dim="wavelength").sel( + t=self.plasma.time_to_calculate) opt_data["xrcs.spectra"]["error"] = np.sqrt( opt_data["xrcs.spectra"] + background.std(dim="wavelength") ** 2 ) @@ -334,10 +345,10 @@ def setup_optimiser(self, model_kwargs): ) self.start_points = self._sample_start_points(sample_high_density=self.sample_high_density) - def _sample_start_points(self, sample_high_density: bool=True): + def _sample_start_points(self, sample_high_density: bool = True): if sample_high_density: start_points = self.bayesmodel.sample_from_high_density_region(self.param_names, self.sampler, - self.nwalkers) + self.nwalkers) else: start_points = self.bayesmodel.sample_from_priors( self.param_names, size=self.nwalkers @@ -367,7 +378,7 @@ def __call__(self, filepath="./results/test/", **kwargs): if __name__ == "__main__": run = BayesWorkflowExample( - pulse=10009, + pulse=None, pulse_to_write=23000101, run="RUN01", diagnostics=["xrcs", "efit", "smmh1", "cxff_pi"], @@ -375,7 +386,7 @@ def __call__(self, filepath="./results/test/", **kwargs): param_names=OPTIMISED_PARAMS, profile_params=DEFAULT_PROFILE_PARAMS, priors=DEFAULT_PRIORS, - model_kwargs={"xrcs_moment_analysis":False, }, + model_kwargs={"xrcs_moment_analysis": False, }, phantoms=True, iterations=200, From e04fe9beaf8dc0c832b51d59ccfab4c49f7f1165 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 8 Aug 2023 14:58:19 +0100 Subject: [PATCH 113/160] formatting --- tests/unit/workflows/test_bayes_workflow_example.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/unit/workflows/test_bayes_workflow_example.py b/tests/unit/workflows/test_bayes_workflow_example.py index a81c5cfe..e4033134 100644 --- a/tests/unit/workflows/test_bayes_workflow_example.py +++ b/tests/unit/workflows/test_bayes_workflow_example.py @@ -117,4 +117,3 @@ def test_worklow_has_results_after_run(self): test = TestBayesWorkflowExample() test.setup_class() - print() From 6fcff7534968ff20ead6f525ff48ce0eaac88cbd Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 8 Aug 2023 14:59:28 +0100 Subject: [PATCH 114/160] refactoring names --- tests/integration/models/test_helike_spectroscopy.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/integration/models/test_helike_spectroscopy.py b/tests/integration/models/test_helike_spectroscopy.py index 58440694..770275d6 100644 --- a/tests/integration/models/test_helike_spectroscopy.py +++ b/tests/integration/models/test_helike_spectroscopy.py @@ -1,5 +1,5 @@ import indica.models.helike_spectroscopy as helike -from indica.models.helike_spectroscopy import helike_LOS_example +from indica.models.helike_spectroscopy import helike_transform_example from indica.models.plasma import example_run as example_plasma @@ -8,13 +8,13 @@ def setup_class(self): self.plasma = example_plasma() # using Phantom self.single_time_point = self.plasma.time_to_calculate[1] self.multiple_time_point = self.plasma.time_to_calculate - self.multiple_channel_los_transform = helike_LOS_example(nchannels=3) - self.single_channel_los_transform = helike_LOS_example(nchannels=1) + self.multiple_channel_los_transform = helike_transform_example(nchannels=3) + self.single_channel_los_transform = helike_transform_example(nchannels=1) self.single_channel_los_transform.set_equilibrium(self.plasma.equilibrium) self.multiple_channel_los_transform.set_equilibrium(self.plasma.equilibrium) def setup_method(self): - self.model = helike.Helike_spectroscopy("diagnostic_name") + self.model = helike.HelikeSpectrometer("diagnostic_name") self.model.set_plasma(self.plasma) def test_helike_moment_runs_with_multiple_LOS( From eba8ca71a0172ef2139abada0c6a5648f30e0e4c Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 8 Aug 2023 15:01:02 +0100 Subject: [PATCH 115/160] deleted old bayes_models tests --- tests/integration/test_bayesmodels.py | 70 --------------------------- 1 file changed, 70 deletions(-) delete mode 100644 tests/integration/test_bayesmodels.py diff --git a/tests/integration/test_bayesmodels.py b/tests/integration/test_bayesmodels.py deleted file mode 100644 index b01f1aca..00000000 --- a/tests/integration/test_bayesmodels.py +++ /dev/null @@ -1,70 +0,0 @@ -import emcee -import flatdict - -from indica.bayesmodels import BayesModels -from indica.bayesmodels import get_uniform -from indica.models.helike_spectroscopy import helike_LOS_example -from indica.models.helike_spectroscopy import Helike_spectroscopy -from indica.models.plasma import example_run - - -class TestBayesModels: - def setup_class(self): - self.plasma = example_run() - self.plasma.time_to_calculate = self.plasma.t[1] - self.los_transform = helike_LOS_example(nchannels=1) - self.los_transform.set_equilibrium(self.plasma.equilibrium) - - def test_simple_run_bayesmodels_with_xrcs(self): - xrcs = Helike_spectroscopy( - name="xrcs", - ) - xrcs.plasma = self.plasma - xrcs.set_los_transform(self.los_transform) - - priors = { - "Te_prof.y0": get_uniform(2e3, 5e3), - "Ti_prof.y0": get_uniform(2e3, 8e3), - } - - bckc = {} - bckc = dict(bckc, **{xrcs.name: {**xrcs(calc_spectra=True)}}) - flat_phantom_data = flatdict.FlatDict(bckc, delimiter=".") - flat_phantom_data["xrcs.spectra"] = flat_phantom_data[ - "xrcs.spectra" - ].expand_dims(dim={"t": [self.plasma.time_to_calculate]}) - - bm = BayesModels( - plasma=self.plasma, - data=flat_phantom_data, - diagnostic_models=[xrcs], - quant_to_optimise=["xrcs.spectra"], - priors=priors, - ) - - # Setup Optimiser - param_names = [ - "Te_prof.y0", - "Ti_prof.y0", - ] - - ndim = param_names.__len__() - nwalkers = ndim * 2 - start_points = bm.sample_from_priors(param_names, size=nwalkers) - - move = [emcee.moves.StretchMove()] - sampler = emcee.EnsembleSampler( - nwalkers, - ndim, - log_prob_fn=bm.ln_posterior, - parameter_names=param_names, - moves=move, - kwargs={}, - ) - sampler.run_mcmc(start_points, 10, progress=False) - - -if __name__ == "__main__": - test = TestBayesModels() - test.setup_class() - test.test_simple_run_bayesmodels_with_xrcs() From b6ece959cd5d0cfd6c167fdea3b5039a2bdbb100 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 8 Aug 2023 15:20:09 +0100 Subject: [PATCH 116/160] reformatted _build_bckc for readability --- indica/bayesmodels.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/indica/bayesmodels.py b/indica/bayesmodels.py index 019e18e3..9be07adf 100644 --- a/indica/bayesmodels.py +++ b/indica/bayesmodels.py @@ -1,7 +1,5 @@ from copy import deepcopy import warnings - -import flatdict import numpy as np from scipy.stats import uniform @@ -79,14 +77,14 @@ def _build_bckc(self, params, **kwargs): _nuisance_params = {param_name.replace(model.name+"_", ""): param_value for param_name, param_value in params.items() if model.name in param_name} - _model_kwargs = {kwarg_name.replace(model.name+"_", ""): + _model_settings = {kwarg_name.replace(model.name+"_", ""): kwarg_value for kwarg_name, kwarg_value in kwargs.items() if model.name in kwarg_name} - model_bckc = {model.name: {**model(**{**_nuisance_params, **_model_kwargs})}} - self.bckc = dict(self.bckc, **model_bckc) - - self.bckc = flatdict.FlatDict(self.bckc, delimiter=".") + _model_kwargs = {**_nuisance_params, **_model_settings} # combine dictionaries + _bckc = model(**_model_kwargs) + _model_bckc = {f"{model.name}.{value_name}": value for value_name, value in _bckc.items()} # prepend model name to bckc + self.bckc = dict(self.bckc, **_model_bckc) return def _ln_likelihood(self): From 2bd93bfc8e110290217f03f096118a98cf930b99 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 8 Aug 2023 15:22:23 +0100 Subject: [PATCH 117/160] moved percentage error to class attribute --- indica/bayesmodels.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/indica/bayesmodels.py b/indica/bayesmodels.py index 9be07adf..43c59268 100644 --- a/indica/bayesmodels.py +++ b/indica/bayesmodels.py @@ -45,12 +45,14 @@ def __init__( quant_to_optimise: list = [], priors: dict = {}, diagnostic_models: list = [], + percent_error: float = None, ): self.plasma = plasma self.data = data self.quant_to_optimise = quant_to_optimise self.diagnostic_models = diagnostic_models self.priors = priors + self.percent_error = percent_error for diag_model in self.diagnostic_models: diag_model.plasma = self.plasma @@ -97,7 +99,7 @@ def _ln_likelihood(self): self.data[key].sel(t=self.plasma.time_to_calculate).astype("float128") ) - exp_error = exp_data * 0.10 # Assume percentage error if none given. + exp_error = exp_data * self.percent_error # Assume percentage error if none given. if hasattr(self.data[key], "error"): if ( self.data[key].error != 0 From e46ac8ad2c1d1d4550525933ba44cd61c6efcfec Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 8 Aug 2023 15:42:22 +0100 Subject: [PATCH 118/160] fixing init percent error --- indica/bayesmodels.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/indica/bayesmodels.py b/indica/bayesmodels.py index 43c59268..cdf6bfec 100644 --- a/indica/bayesmodels.py +++ b/indica/bayesmodels.py @@ -45,7 +45,7 @@ def __init__( quant_to_optimise: list = [], priors: dict = {}, diagnostic_models: list = [], - percent_error: float = None, + percent_error: float = 0.10, ): self.plasma = plasma self.data = data From cc0aafd4ebdb9ca0191f0144634acaf82e10f28f Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 8 Aug 2023 15:52:10 +0100 Subject: [PATCH 119/160] black formatting --- indica/bayesmodels.py | 69 ++-- indica/models/charge_exchange.py | 6 +- indica/models/helike_spectroscopy.py | 1 - indica/models/interferometry.py | 3 +- indica/models/plasma.py | 8 +- indica/readers/read_st40.py | 2 - indica/workflows/abstract_bayes_workflow.py | 326 ++++++++++++------ indica/workflows/bayes_plots.py | 16 +- indica/workflows/bayes_workflow_example.py | 143 ++++---- indica/writers/bda_tree.py | 100 +++--- .../workflows/test_bayes_workflow_example.py | 74 ++-- 11 files changed, 463 insertions(+), 285 deletions(-) diff --git a/indica/bayesmodels.py b/indica/bayesmodels.py index cdf6bfec..c8801664 100644 --- a/indica/bayesmodels.py +++ b/indica/bayesmodels.py @@ -7,8 +7,15 @@ warnings.simplefilter("ignore", category=FutureWarning) -PROFILES = ["electron_temperature", "electron_density", "ion_temperature", "ion_density", - "impurity_density", "fast_density", "neutral_density"] +PROFILES = [ + "electron_temperature", + "electron_density", + "ion_temperature", + "ion_density", + "impurity_density", + "fast_density", + "neutral_density", +] def gaussian(x, mean, sigma): @@ -76,16 +83,26 @@ def _build_bckc(self, params, **kwargs): self.bckc: dict = {} for model in self.diagnostic_models: # removes "model.name_" from params and kwargs then passes them to model e.g. xrcs_background -> background - _nuisance_params = {param_name.replace(model.name+"_", ""): - param_value for param_name, param_value in params.items() - if model.name in param_name} - _model_settings = {kwarg_name.replace(model.name+"_", ""): - kwarg_value for kwarg_name, kwarg_value in kwargs.items() - if model.name in kwarg_name} - - _model_kwargs = {**_nuisance_params, **_model_settings} # combine dictionaries + _nuisance_params = { + param_name.replace(model.name + "_", ""): param_value + for param_name, param_value in params.items() + if model.name in param_name + } + _model_settings = { + kwarg_name.replace(model.name + "_", ""): kwarg_value + for kwarg_name, kwarg_value in kwargs.items() + if model.name in kwarg_name + } + + _model_kwargs = { + **_nuisance_params, + **_model_settings, + } # combine dictionaries _bckc = model(**_model_kwargs) - _model_bckc = {f"{model.name}.{value_name}": value for value_name, value in _bckc.items()} # prepend model name to bckc + _model_bckc = { + f"{model.name}.{value_name}": value + for value_name, value in _bckc.items() + } # prepend model name to bckc self.bckc = dict(self.bckc, **_model_bckc) return @@ -99,12 +116,16 @@ def _ln_likelihood(self): self.data[key].sel(t=self.plasma.time_to_calculate).astype("float128") ) - exp_error = exp_data * self.percent_error # Assume percentage error if none given. + exp_error = ( + exp_data * self.percent_error + ) # Assume percentage error if none given. if hasattr(self.data[key], "error"): if ( self.data[key].error != 0 ).any(): # TODO: Some models have an error of 0 given - exp_error = self.data[key].error.sel(t=self.plasma.time_to_calculate) + exp_error = self.data[key].error.sel( + t=self.plasma.time_to_calculate + ) _ln_likelihood = np.log(gaussian(model_data, exp_data, exp_error)) # treat channel as key dim which isn't averaged like other dims @@ -171,10 +192,10 @@ def sample_from_priors(self, param_names, size=10): samples = samples[:, 0:size] return samples.transpose() - def sample_from_high_density_region(self, param_names: list, sampler: object, nwalkers: int, nsamples = 100): - start_points = self.sample_from_priors( - param_names, size=nsamples - ) + def sample_from_high_density_region( + self, param_names: list, sampler: object, nwalkers: int, nsamples=100 + ): + start_points = self.sample_from_priors(param_names, size=nsamples) ln_prob, _ = sampler.compute_log_prob(start_points) num_best_points = int(nsamples * 0.05) @@ -190,17 +211,13 @@ def sample_from_high_density_region(self, param_names: list, sampler: object, nw best_points_std * 2, size=(nwalkers * 5, len(param_names)), ) - start = { - name: sample[:, idx] for idx, name in enumerate(param_names) - } + start = {name: sample[:, idx] for idx, name in enumerate(param_names)} ln_prior = self._ln_prior(start) # Convert from dictionary of arrays -> array, # then filtering out where ln_prior is -infinity - accepted_samples = np.array(list(start.values()))[ - :, ln_prior != -np.inf - ] + accepted_samples = np.array(list(start.values()))[:, ln_prior != -np.inf] samples = np.append(samples, accepted_samples, axis=1) - start_points = samples[:, 0: nwalkers].transpose() + start_points = samples[:, 0:nwalkers].transpose() return start_points def ln_posterior(self, parameters: dict, **kwargs): @@ -234,7 +251,9 @@ def ln_posterior(self, parameters: dict, **kwargs): kin_profs = {} for profile_key in PROFILES: if hasattr(self.plasma, profile_key): - kin_profs[profile_key] = getattr(self.plasma, profile_key).sel(t=self.plasma.time_to_calculate) + kin_profs[profile_key] = getattr(self.plasma, profile_key).sel( + t=self.plasma.time_to_calculate + ) else: raise ValueError(f"plasma does not have attribute {profile_key}") diff --git a/indica/models/charge_exchange.py b/indica/models/charge_exchange.py index 1aae94ac..1f0eeda8 100644 --- a/indica/models/charge_exchange.py +++ b/indica/models/charge_exchange.py @@ -22,7 +22,6 @@ def __init__( element: str = "c", instrument_method="get_charge_exchange", ): - self.name = name self.element = element self.instrument_method = instrument_method @@ -106,7 +105,8 @@ def __call__( return self.bckc -def pi_transform_example(nchannels:int): + +def pi_transform_example(nchannels: int): x_positions = np.linspace(0.2, 0.8, nchannels) y_positions = np.linspace(0.0, 0.0, nchannels) z_positions = np.linspace(0.0, 0.0, nchannels) @@ -120,13 +120,13 @@ def pi_transform_example(nchannels:int): ) return transect_transform + def example_run( pulse: int = None, diagnostic_name: str = "cxrs", plasma=None, plot=False, ): - # TODO: LOS sometimes crossing bad EFIT reconstruction if plasma is None: diff --git a/indica/models/helike_spectroscopy.py b/indica/models/helike_spectroscopy.py index d4538b20..12642af3 100644 --- a/indica/models/helike_spectroscopy.py +++ b/indica/models/helike_spectroscopy.py @@ -686,7 +686,6 @@ def __call__( self.intensity = self._calculate_intensity() self.spectra = self._make_spectra() - self._calculate_los_integral( calc_rho=calc_rho, ) diff --git a/indica/models/interferometry.py b/indica/models/interferometry.py index 29a6b42b..49cf391b 100644 --- a/indica/models/interferometry.py +++ b/indica/models/interferometry.py @@ -24,7 +24,6 @@ def __init__( name: str, instrument_method="get_interferometry", ): - self.name = name self.instrument_method = instrument_method self.quantities = AVAILABLE_QUANTITIES[self.instrument_method] @@ -88,6 +87,7 @@ def __call__( return self.bckc + def smmh1_transform_example(): los_start = np.array([[0.8, 0, 0], [0.8, 0, -0.1], [0.8, 0, -0.2]]) los_end = np.array([[0.17, 0, 0], [0.17, 0, -0.25], [0.17, 0, -0.2]]) @@ -106,6 +106,7 @@ def smmh1_transform_example(): ) return los_transform + def example_run(pulse: int = None, plasma=None, plot=False): if plasma is None: plasma = example_plasma(pulse=pulse) diff --git a/indica/models/plasma.py b/indica/models/plasma.py index 4ce7bac6..1715dd3f 100644 --- a/indica/models/plasma.py +++ b/indica/models/plasma.py @@ -462,7 +462,9 @@ def assign_profiles(self, impurity_to_profile=None): self.electron_density.loc[dict(t=t)] = self.Ne_prof() self.calc_impurity_density(self.impurities) self.electron_temperature.loc[dict(t=t)] = self.Te_prof() - self.ion_temperature.loc[dict(t=t)] = self.Ti_prof(y0_ref=self.Te_prof.y0) # Temp default Ti_ref + self.ion_temperature.loc[dict(t=t)] = self.Ti_prof( + y0_ref=self.Te_prof.y0 + ) # Temp default Ti_ref self.toroidal_rotation.loc[dict(t=t)] = self.Vrot_prof() self.neutral_density.loc[dict(t=t)] = self.Nh_prof() if ( @@ -1281,7 +1283,9 @@ def example_run( plasma.Nimp_prof.peaking = nimp_peaking[i] plasma.Nimp_prof.y0 = nimp_y0[i] plasma.Nimp_prof.wcenter = nimp_wcenter[i] - plasma.assign_profiles(impurity_to_profile="ar",) + plasma.assign_profiles( + impurity_to_profile="ar", + ) if pulse is None: equilibrium_data = fake_equilibrium_data( diff --git a/indica/readers/read_st40.py b/indica/readers/read_st40.py index 07e95aec..b23e25bd 100644 --- a/indica/readers/read_st40.py +++ b/indica/readers/read_st40.py @@ -153,7 +153,6 @@ def map_diagnostics(self, instruments: list, map_raw: bool = False): break def filter_data(self, instruments: list): - if not hasattr(self, "binned_data"): raise ValueError( "Bin data before filtering. No action permitted on raw data structure!" @@ -267,7 +266,6 @@ def __call__( map_diagnostics: bool = False, debug: bool = False, ): - if instruments is None: instruments = list(REVISIONS.keys()) if revisions is None: diff --git a/indica/workflows/abstract_bayes_workflow.py b/indica/workflows/abstract_bayes_workflow.py index 215742ff..86057daa 100644 --- a/indica/workflows/abstract_bayes_workflow.py +++ b/indica/workflows/abstract_bayes_workflow.py @@ -8,16 +8,17 @@ from indica.readers.read_st40 import ReadST40 from indica.equilibrium import fake_equilibrium + class AbstractBayesWorkflow(ABC): @abstractmethod - def __init__(self, - phantoms = None, - diagnostics = None, - param_names=None, - opt_quantity=None, - priors=None, - ): - + def __init__( + self, + phantoms=None, + diagnostics=None, + param_names=None, + opt_quantity=None, + priors=None, + ): self.phantoms = phantoms self.diagnostics = diagnostics self.param_names = param_names @@ -31,10 +32,14 @@ def __init__(self, self.setup_opt_data(self.phantoms) self.setup_optimiser() - def read_test_data(self, diagnostic_transforms:dict): + def read_test_data(self, diagnostic_transforms: dict): # Used with phantom data for purposes of tests self.reader = None - self.equilibrium = fake_equilibrium(self.tstart, self.tend, self.dt, ) + self.equilibrium = fake_equilibrium( + self.tstart, + self.tend, + self.dt, + ) self.transforms = diagnostic_transforms self.data = {} @@ -92,21 +97,34 @@ def setup_optimiser(self, model_kwargs): def save_phantom_profiles(self, kinetic_profiles=None): if kinetic_profiles is None: - kinetic_profiles = ["electron_density", "impurity_density", "electron_temperature", - "ion_temperature", "ion_density", "fast_density", "neutral_density"] + kinetic_profiles = [ + "electron_density", + "impurity_density", + "electron_temperature", + "ion_temperature", + "ion_density", + "fast_density", + "neutral_density", + ] if self.phantoms: - phantom_profiles = {profile_key: getattr(self.plasma, profile_key).sel( - t=self.plasma.time_to_calculate).copy() - for profile_key in kinetic_profiles} + phantom_profiles = { + profile_key: getattr(self.plasma, profile_key) + .sel(t=self.plasma.time_to_calculate) + .copy() + for profile_key in kinetic_profiles + } else: - phantom_profiles = {profile_key: getattr(self.plasma, profile_key).sel( - t=self.plasma.time_to_calculate) * 0 - for profile_key in kinetic_profiles} + phantom_profiles = { + profile_key: getattr(self.plasma, profile_key).sel( + t=self.plasma.time_to_calculate + ) + * 0 + for profile_key in kinetic_profiles + } self.phantom_profiles = phantom_profiles def _build_result_dict(self): - """ Returns @@ -132,89 +150,134 @@ def _build_result_dict(self): "BURN_FRAC": self.burn_frac, "ITER": self.iterations, "NWALKERS": self.nwalkers, - "MODEL_KWARGS": self.model_kwargs, "OPT_QUANTITY": self.opt_quantity, "PARAM_NAMES": self.param_names, - "PULSE": self.pulse, "DT": self.dt, - "IMPURITIES":self.plasma.impurities, + "IMPURITIES": self.plasma.impurities, "MAIN_ION": self.plasma.main_ion, } result["INPUT"]["WORKFLOW"] = { - diag_name.upper(): - {"PULSE": self.pulse, # Change this if different pulses used for diagnostics - "USAGE": "".join([quantity[1] for quantity in quant_list if quantity[0] == diag_name]), - "RUN": "PLACEHOLDER", - } + diag_name.upper(): { + "PULSE": self.pulse, # Change this if different pulses used for diagnostics + "USAGE": "".join( + [quantity[1] for quantity in quant_list if quantity[0] == diag_name] + ), + "RUN": "PLACEHOLDER", + } for diag_name in self.diagnostics } - result["MODEL_DATA"] = {diag_name.upper(): - {quantity[1].upper(): self.blobs[f"{quantity[0]}.{quantity[1]}"] - for quantity in quant_list if quantity[0] == diag_name} - for diag_name in self.diagnostics - } + result["MODEL_DATA"] = { + diag_name.upper(): { + quantity[1].upper(): self.blobs[f"{quantity[0]}.{quantity[1]}"] + for quantity in quant_list + if quantity[0] == diag_name + } + for diag_name in self.diagnostics + } result["MODEL_DATA"]["SAMPLES"] = self.samples - result["DIAG_DATA"] = {diag_name.upper(): - {quantity[1].upper(): self.opt_data[f"{quantity[0]}.{quantity[1]}"] - for quantity in quant_list if quantity[0] == diag_name} - for diag_name in self.diagnostics - } - + result["DIAG_DATA"] = { + diag_name.upper(): { + quantity[1].upper(): self.opt_data[f"{quantity[0]}.{quantity[1]}"] + for quantity in quant_list + if quantity[0] == diag_name + } + for diag_name in self.diagnostics + } result["PHANTOMS"] = { "FLAG": self.phantoms, "NE": self.phantom_profiles["electron_density"], "TE": self.phantom_profiles["electron_temperature"], - "TI": self.phantom_profiles["ion_temperature"].sel(element=self.plasma.main_ion), - "NI": self.phantom_profiles["ion_density"].sel(element=self.plasma.main_ion), + "TI": self.phantom_profiles["ion_temperature"].sel( + element=self.plasma.main_ion + ), + "NI": self.phantom_profiles["ion_density"].sel( + element=self.plasma.main_ion + ), "NNEUTR": self.phantom_profiles["neutral_density"], "NFAST": self.phantom_profiles["fast_density"], } - result["PHANTOMS"].update({ - f"NIZ{num_imp+1}": self.phantom_profiles["impurity_density"].sel(element=imp) - for num_imp, imp in enumerate(self.plasma.impurities) - }) - result["PHANTOMS"].update({ - f"TIZ{num_imp + 1}": self.phantom_profiles["ion_temperature"].sel(element=imp) - for num_imp, imp in enumerate(self.plasma.impurities) - }) + result["PHANTOMS"].update( + { + f"NIZ{num_imp+1}": self.phantom_profiles["impurity_density"].sel( + element=imp + ) + for num_imp, imp in enumerate(self.plasma.impurities) + } + ) + result["PHANTOMS"].update( + { + f"TIZ{num_imp + 1}": self.phantom_profiles["ion_temperature"].sel( + element=imp + ) + for num_imp, imp in enumerate(self.plasma.impurities) + } + ) result["PROFILES"] = { "RHO_POLOIDAL": self.plasma.rho, "RHO_TOR": self.plasma.equilibrium.rhotor.interp(t=self.plasma.t), - "NE": self.blobs["electron_density"].median(dim="index"), - "NI": self.blobs["ion_density"].sel(element=self.plasma.main_ion).median(dim="index"), + "NI": self.blobs["ion_density"] + .sel(element=self.plasma.main_ion) + .median(dim="index"), "TE": self.blobs["electron_temperature"].median(dim="index"), - "TI": self.blobs["ion_temperature"].sel(element=self.plasma.main_ion).median(dim="index"), + "TI": self.blobs["ion_temperature"] + .sel(element=self.plasma.main_ion) + .median(dim="index"), "NFAST": self.blobs["fast_density"].median(dim="index"), "NNEUTR": self.blobs["neutral_density"].median(dim="index"), - "NE_ERR": self.blobs["electron_density"].std(dim="index"), - "NI_ERR": self.blobs["ion_density"].sel(element=self.plasma.main_ion).std(dim="index"), + "NI_ERR": self.blobs["ion_density"] + .sel(element=self.plasma.main_ion) + .std(dim="index"), "TE_ERR": self.blobs["electron_temperature"].std(dim="index"), - "TI_ERR": self.blobs["ion_temperature"].sel(element=self.plasma.main_ion).std(dim="index"), + "TI_ERR": self.blobs["ion_temperature"] + .sel(element=self.plasma.main_ion) + .std(dim="index"), "NFAST_ERR": self.blobs["fast_density"].std(dim="index"), "NNEUTR_ERR": self.blobs["neutral_density"].std(dim="index"), - } - result["PROFILES"] = {**result["PROFILES"], **{ - f"NIZ{num_imp+1}": self.blobs["impurity_density"].sel(element=imp).median(dim="index") - for num_imp, imp in enumerate(self.plasma.impurities)}} - result["PROFILES"] = {**result["PROFILES"], **{ - f"NIZ{num_imp+1}_ERR": self.blobs["impurity_density"].sel(element=imp).std(dim="index") - for num_imp, imp in enumerate(self.plasma.impurities)}} - result["PROFILES"] = {**result["PROFILES"], **{ - f"TIZ{num_imp + 1}": self.blobs["ion_temperature"].sel(element=imp).median(dim="index") - for num_imp, imp in enumerate(self.plasma.impurities)}} - result["PROFILES"] = {**result["PROFILES"], **{ - f"TIZ{num_imp + 1}_ERR": self.blobs["ion_temperature"].sel(element=imp).std(dim="index") - for num_imp, imp in enumerate(self.plasma.impurities)}} - + result["PROFILES"] = { + **result["PROFILES"], + **{ + f"NIZ{num_imp+1}": self.blobs["impurity_density"] + .sel(element=imp) + .median(dim="index") + for num_imp, imp in enumerate(self.plasma.impurities) + }, + } + result["PROFILES"] = { + **result["PROFILES"], + **{ + f"NIZ{num_imp+1}_ERR": self.blobs["impurity_density"] + .sel(element=imp) + .std(dim="index") + for num_imp, imp in enumerate(self.plasma.impurities) + }, + } + result["PROFILES"] = { + **result["PROFILES"], + **{ + f"TIZ{num_imp + 1}": self.blobs["ion_temperature"] + .sel(element=imp) + .median(dim="index") + for num_imp, imp in enumerate(self.plasma.impurities) + }, + } + result["PROFILES"] = { + **result["PROFILES"], + **{ + f"TIZ{num_imp + 1}_ERR": self.blobs["ion_temperature"] + .sel(element=imp) + .std(dim="index") + for num_imp, imp in enumerate(self.plasma.impurities) + }, + } result["PROFILE_STAT"] = { "SAMPLES": self.samples, @@ -226,13 +289,20 @@ def _build_result_dict(self): "NFAST": self.blobs["fast_density"], "NNEUTR": self.blobs["neutral_density"], } - result["PROFILE_STAT"] = {**result["PROFILE_STAT"], **{ - f"NIZ{num_imp+1}": self.blobs["impurity_density"].sel(element=imp) - for num_imp, imp in enumerate(self.plasma.impurities)}} - result["PROFILE_STAT"] = {**result["PROFILE_STAT"], **{ - f"TIZ{num_imp + 1}": self.blobs["ion_temperature"].sel(element=imp) - for num_imp, imp in enumerate(self.plasma.impurities)}} - + result["PROFILE_STAT"] = { + **result["PROFILE_STAT"], + **{ + f"NIZ{num_imp+1}": self.blobs["impurity_density"].sel(element=imp) + for num_imp, imp in enumerate(self.plasma.impurities) + }, + } + result["PROFILE_STAT"] = { + **result["PROFILE_STAT"], + **{ + f"TIZ{num_imp + 1}": self.blobs["ion_temperature"].sel(element=imp) + for num_imp, imp in enumerate(self.plasma.impurities) + }, + } result["OPTIMISATION"] = { "ACCEPT_FRAC": self.accept_frac, @@ -242,37 +312,80 @@ def _build_result_dict(self): } result["GLOBAL"] = { - "TI0": self.blobs["ion_temperature"].sel(element=self.plasma.main_ion).sel(rho_poloidal=0, method="nearest").median(dim="index"), - "TE0": self.blobs["electron_temperature"].sel(rho_poloidal=0, method="nearest").median(dim="index"), - "NE0": self.blobs["electron_density"].sel(rho_poloidal=0, method="nearest").median(dim="index"), - "NI0": self.blobs["ion_density"].sel(element=self.plasma.main_ion).sel(rho_poloidal=0, method="nearest").median(dim="index"), - "TI0_ERR": self.blobs["ion_temperature"].sel(element=self.plasma.main_ion).sel(rho_poloidal=0, method="nearest").std(dim="index"), - "TE0_ERR": self.blobs["electron_temperature"].sel(rho_poloidal=0, method="nearest").std(dim="index"), - "NE0_ERR": self.blobs["electron_density"].sel(rho_poloidal=0, method="nearest").std(dim="index"), - "NI0_ERR": self.blobs["ion_density"].sel(element=self.plasma.main_ion).sel(rho_poloidal=0, method="nearest").std(dim="index"), + "TI0": self.blobs["ion_temperature"] + .sel(element=self.plasma.main_ion) + .sel(rho_poloidal=0, method="nearest") + .median(dim="index"), + "TE0": self.blobs["electron_temperature"] + .sel(rho_poloidal=0, method="nearest") + .median(dim="index"), + "NE0": self.blobs["electron_density"] + .sel(rho_poloidal=0, method="nearest") + .median(dim="index"), + "NI0": self.blobs["ion_density"] + .sel(element=self.plasma.main_ion) + .sel(rho_poloidal=0, method="nearest") + .median(dim="index"), + "TI0_ERR": self.blobs["ion_temperature"] + .sel(element=self.plasma.main_ion) + .sel(rho_poloidal=0, method="nearest") + .std(dim="index"), + "TE0_ERR": self.blobs["electron_temperature"] + .sel(rho_poloidal=0, method="nearest") + .std(dim="index"), + "NE0_ERR": self.blobs["electron_density"] + .sel(rho_poloidal=0, method="nearest") + .std(dim="index"), + "NI0_ERR": self.blobs["ion_density"] + .sel(element=self.plasma.main_ion) + .sel(rho_poloidal=0, method="nearest") + .std(dim="index"), + } + result["GLOBAL"] = { + **result["GLOBAL"], + **{ + f"TI0Z{num_imp + 1}": self.blobs["ion_temperature"] + .sel(element=imp) + .sel(rho_poloidal=0, method="nearest") + .median(dim="index") + for num_imp, imp in enumerate(self.plasma.impurities) + }, + } + result["GLOBAL"] = { + **result["GLOBAL"], + **{ + f"TI0Z{num_imp + 1}_ERR": self.blobs["ion_temperature"] + .sel(element=imp) + .sel(rho_poloidal=0, method="nearest") + .std(dim="index") + for num_imp, imp in enumerate(self.plasma.impurities) + }, + } + result["GLOBAL"] = { + **result["GLOBAL"], + **{ + f"NI0Z{num_imp + 1}": self.blobs["impurity_density"] + .sel(element=imp) + .sel(rho_poloidal=0, method="nearest") + .median(dim="index") + for num_imp, imp in enumerate(self.plasma.impurities) + }, + } + result["GLOBAL"] = { + **result["GLOBAL"], + **{ + f"NI0Z{num_imp + 1}_ERR": self.blobs["impurity_density"] + .sel(element=imp) + .sel(rho_poloidal=0, method="nearest") + .std(dim="index") + for num_imp, imp in enumerate(self.plasma.impurities) + }, } - result["GLOBAL"] = {**result["GLOBAL"], **{ - f"TI0Z{num_imp + 1}": self.blobs["ion_temperature"].sel(element=imp).sel(rho_poloidal=0, method="nearest").median(dim="index") - for num_imp, imp in enumerate(self.plasma.impurities)}} - result["GLOBAL"] = {**result["GLOBAL"], **{ - f"TI0Z{num_imp + 1}_ERR": self.blobs["ion_temperature"].sel(element=imp).sel(rho_poloidal=0, - method="nearest").std( - dim="index") - for num_imp, imp in enumerate(self.plasma.impurities)}} - result["GLOBAL"] = {**result["GLOBAL"], **{ - f"NI0Z{num_imp + 1}": self.blobs["impurity_density"].sel(element=imp).sel(rho_poloidal=0, method="nearest").median(dim="index") - for num_imp, imp in enumerate(self.plasma.impurities)}} - result["GLOBAL"] = {**result["GLOBAL"], **{ - f"NI0Z{num_imp + 1}_ERR": self.blobs["impurity_density"].sel(element=imp).sel(rho_poloidal=0, - method="nearest").std( - dim="index") - for num_imp, imp in enumerate(self.plasma.impurities)}} self.result = result return self.result def run_sampler(self): - """ TODO: unsure if keeping in abstract class is best practice @@ -284,7 +397,11 @@ def run_sampler(self): """ self.autocorr = sample_with_autocorr( - self.sampler, self.start_points, self.iterations, self.param_names.__len__(), auto_sample=10 + self.sampler, + self.start_points, + self.iterations, + self.param_names.__len__(), + auto_sample=10, ) blobs = self.sampler.get_blobs( discard=int(self.iterations * self.burn_frac), flat=True @@ -306,7 +423,6 @@ def run_sampler(self): self.post_sample = self.sampler.get_chain(flat=True) self.result = self._build_result_dict() - def save_pickle(self, filepath): if filepath: Path(filepath).mkdir(parents=True, exist_ok=True) @@ -345,11 +461,11 @@ def sample_with_autocorr(sampler, start_points, iterations, n_params, auto_sampl if sampler.iteration % auto_sample: continue new_tau = sampler.get_autocorr_time(tol=0) - autocorr[sampler.iteration - 1, ] = new_tau + autocorr[sampler.iteration - 1,] = new_tau converged = np.all(new_tau * 50 < sampler.iteration) converged &= np.all(np.abs(old_tau - new_tau) / new_tau < 0.01) if converged: break old_tau = new_tau - autocorr = autocorr[: sampler.iteration, ] - return autocorr \ No newline at end of file + autocorr = autocorr[: sampler.iteration,] + return autocorr diff --git a/indica/workflows/bayes_plots.py b/indica/workflows/bayes_plots.py index a4e06773..8695969d 100644 --- a/indica/workflows/bayes_plots.py +++ b/indica/workflows/bayes_plots.py @@ -176,7 +176,7 @@ def violinplot( axs.set_xlabel(key) top = axs.get_ylim()[1] bot = axs.get_ylim()[0] - axs.set_ylim( top=top * 1.1, bottom=bot * 0.9) + axs.set_ylim(top=top * 1.1, bottom=bot * 0.9) axs.set_ylabel(f"{ylabel}") y = diag_data[key].sel(t=data[key].t).values axs.errorbar( @@ -218,17 +218,14 @@ def histograms(data, diag_data, filename): plt.close() - def plot_autocorr(autocorr, param_names, figheader, filetype=".png"): - plt.figure() - x_data = np.ones(shape=(autocorr.shape)) * np.arange(0, autocorr[:,0].__len__())[:,None] - plt.plot( - np.where(np.isfinite(autocorr), x_data, np.nan), - autocorr, - "x" + x_data = ( + np.ones(shape=(autocorr.shape)) + * np.arange(0, autocorr[:, 0].__len__())[:, None] ) + plt.plot(np.where(np.isfinite(autocorr), x_data, np.nan), autocorr, "x") plt.legend(param_names) plt.xlabel("iterations") plt.ylabel("auto-correlation time (iterations)") @@ -242,7 +239,6 @@ def plot_bayes_result( filetype=".png", **kwargs, ): - # delete all but pickle in directory for root, dirs, files in os.walk(figheader): for f in files: @@ -380,7 +376,7 @@ def plot_bayes_result( filetype=filetype, phantom_profile=phantom_profiles, color="blue", - sharefig=True + sharefig=True, ) key = "NI" plot_profile( diff --git a/indica/workflows/bayes_workflow_example.py b/indica/workflows/bayes_workflow_example.py index b2d5f4e3..02b20ec1 100644 --- a/indica/workflows/bayes_workflow_example.py +++ b/indica/workflows/bayes_workflow_example.py @@ -9,7 +9,10 @@ from indica.writers.bda_tree import create_nodes, write_nodes, check_analysis_run from indica.models.interferometry import Interferometry, smmh1_transform_example -from indica.models.helike_spectroscopy import HelikeSpectrometer, helike_transform_example +from indica.models.helike_spectroscopy import ( + HelikeSpectrometer, + helike_transform_example, +) from indica.models.charge_exchange import ChargeExchange, pi_transform_example from indica.models.equilibrium_reconstruction import EquilibriumReconstruction from indica.models.plasma import Plasma @@ -102,29 +105,27 @@ class BayesWorkflowExample(AbstractBayesWorkflow): def __init__( - self, - pulse: int = None, - pulse_to_write: bool = None, - run: str = "RUN01", - diagnostics: list = None, - param_names: list = None, - opt_quantity: list = None, - priors: dict = None, - profile_params: dict = None, - phantoms: bool = False, - model_kwargs: dict = {}, - - nwalkers: int = 50, - tstart: float = 0.02, - tend: float = 0.10, - dt: float = 0.01, - tsample: float = 0.06, - iterations: int = 100, - burn_frac: float = 0, - - mds_write: bool = False, - plot: bool = True, - sample_high_density: bool = False, + self, + pulse: int = None, + pulse_to_write: bool = None, + run: str = "RUN01", + diagnostics: list = None, + param_names: list = None, + opt_quantity: list = None, + priors: dict = None, + profile_params: dict = None, + phantoms: bool = False, + model_kwargs: dict = {}, + nwalkers: int = 50, + tstart: float = 0.02, + tend: float = 0.10, + dt: float = 0.01, + tsample: float = 0.06, + iterations: int = 100, + burn_frac: float = 0, + mds_write: bool = False, + plot: bool = True, + sample_high_density: bool = False, ): self.pulse = pulse self.pulse_to_write = pulse_to_write @@ -149,18 +150,30 @@ def __init__( self.plot = plot self.sample_high_density = sample_high_density - for attribute in ["param_names", "opt_quantity", "priors", "diagnostics", "profile_params"]: + for attribute in [ + "param_names", + "opt_quantity", + "priors", + "diagnostics", + "profile_params", + ]: if getattr(self, attribute) is None: raise ValueError(f"{attribute} needs to be defined") if self.pulse is None and self.phantoms is False: - raise ValueError("Set phantoms to True when running phantom plasma i.e. pulse=None") + raise ValueError( + "Set phantoms to True when running phantom plasma i.e. pulse=None" + ) if pulse is None: print("Running in test mode") - self.read_test_data({"xrcs": helike_transform_example(1), - "smmh1": smmh1_transform_example(), - "cxff_pi": pi_transform_example(5)}) + self.read_test_data( + { + "xrcs": helike_transform_example(1), + "smmh1": smmh1_transform_example(), + "cxff_pi": pi_transform_example(5), + } + ) else: self.read_data(self.diagnostics) self.setup_plasma() @@ -220,12 +233,14 @@ def setup_models(self, diagnostics: list): window_vector = None if hasattr(self, "data"): if diag in self.data.keys(): - window_vector = self.data[diag]["spectra"].wavelength.values * 0.1 + window_vector = ( + self.data[diag]["spectra"].wavelength.values * 0.1 + ) model = HelikeSpectrometer( name="xrcs", window_masks=[slice(0.394, 0.396)], - window_vector=window_vector + window_vector=window_vector, ) model.set_los_transform(los_transform) @@ -253,46 +268,42 @@ def _phantom_data(self, noise=False, noise_factor=0.1): if "smmh1" in self.diagnostics: opt_data["smmh1.ne"] = ( self.models["smmh1"]() - .pop("ne") - .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) + .pop("ne") + .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) ) if "xrcs" in self.diagnostics: opt_data["xrcs.spectra"] = ( self.models["xrcs"]() - .pop("spectra") - .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) - ) - opt_data["xrcs.spectra"]["error"] = np.sqrt( - opt_data["xrcs.spectra"] + .pop("spectra") + .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) ) + opt_data["xrcs.spectra"]["error"] = np.sqrt(opt_data["xrcs.spectra"]) if "cxff_pi" in self.diagnostics: cxrs_data = ( self.models["cxff_pi"]() - .pop("ti") - .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) + .pop("ti") + .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) ) opt_data["cxff_pi.ti"] = cxrs_data.where(cxrs_data.channel == 2, drop=True) if "efit" in self.diagnostics: opt_data["efit.wp"] = ( self.models["efit"]() - .pop("wp") - .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) + .pop("wp") + .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) ) if noise: opt_data["smmh1.ne"] = opt_data["smmh1.ne"] + self.opt_data[ "smmh1.ne" ].max().values * np.random.normal(0, noise_factor, None) - opt_data["xrcs.spectra"] = opt_data[ - "xrcs.spectra" - ] + np.random.normal( + opt_data["xrcs.spectra"] = opt_data["xrcs.spectra"] + np.random.normal( 0, np.sqrt(opt_data["xrcs.spectra"].values[0,]), opt_data["xrcs.spectra"].shape[1], ) - opt_data["cxff_pi.ti"] = opt_data[ - "cxff_pi.ti" - ] + opt_data["cxff_pi.ti"].max().values * np.random.normal( + opt_data["cxff_pi.ti"] = opt_data["cxff_pi.ti"] + opt_data[ + "cxff_pi.ti" + ].max().values * np.random.normal( 0, noise_factor, opt_data["cxff_pi.ti"].shape[1] ) opt_data["efit.wp"] = opt_data["efit.wp"] + opt_data[ @@ -304,15 +315,16 @@ def _exp_data(self): opt_data = flatdict.FlatDict(self.data, ".") if "xrcs" in self.diagnostics: opt_data["xrcs.spectra"]["wavelength"] = ( - opt_data["xrcs.spectra"].wavelength * 0.1 + opt_data["xrcs.spectra"].wavelength * 0.1 ) background = opt_data["xrcs.spectra"].where( (opt_data["xrcs.spectra"].wavelength < 0.392) & (opt_data["xrcs.spectra"].wavelength > 0.388), drop=True, ) - self.model_kwargs["xrcs_background"] = background.mean(dim="wavelength").sel( - t=self.plasma.time_to_calculate) + self.model_kwargs["xrcs_background"] = background.mean( + dim="wavelength" + ).sel(t=self.plasma.time_to_calculate) opt_data["xrcs.spectra"]["error"] = np.sqrt( opt_data["xrcs.spectra"] + background.std(dim="wavelength") ** 2 ) @@ -324,7 +336,6 @@ def _exp_data(self): return opt_data def setup_optimiser(self, model_kwargs): - self.bayesmodel = BayesModels( plasma=self.plasma, data=self.opt_data, @@ -343,12 +354,15 @@ def setup_optimiser(self, model_kwargs): moves=self.move, kwargs=model_kwargs, ) - self.start_points = self._sample_start_points(sample_high_density=self.sample_high_density) + self.start_points = self._sample_start_points( + sample_high_density=self.sample_high_density + ) def _sample_start_points(self, sample_high_density: bool = True): if sample_high_density: - start_points = self.bayesmodel.sample_from_high_density_region(self.param_names, self.sampler, - self.nwalkers) + start_points = self.bayesmodel.sample_from_high_density_region( + self.param_names, self.sampler, self.nwalkers + ) else: start_points = self.bayesmodel.sample_from_priors( self.param_names, size=self.nwalkers @@ -356,12 +370,13 @@ def _sample_start_points(self, sample_high_density: bool = True): return start_points def __call__(self, filepath="./results/test/", **kwargs): - if self.mds_write: # check_analysis_run(self.pulse, self.run) - self.node_structure = create_nodes(pulse_to_write=self.pulse_to_write, - diagnostic_quantities=self.opt_quantity, - mode="NEW") + self.node_structure = create_nodes( + pulse_to_write=self.pulse_to_write, + diagnostic_quantities=self.opt_quantity, + mode="NEW", + ) self.run_sampler() self.save_pickle(filepath=filepath) @@ -386,17 +401,19 @@ def __call__(self, filepath="./results/test/", **kwargs): param_names=OPTIMISED_PARAMS, profile_params=DEFAULT_PROFILE_PARAMS, priors=DEFAULT_PRIORS, - model_kwargs={"xrcs_moment_analysis": False, }, + model_kwargs={ + "xrcs_moment_analysis": False, + }, phantoms=True, - iterations=200, nwalkers=50, burn_frac=0.10, dt=0.005, tsample=0.060, - mds_write=True, plot=True, sample_high_density=False, ) - results = run(filepath="./results/test/", ) + results = run( + filepath="./results/test/", + ) diff --git a/indica/writers/bda_tree.py b/indica/writers/bda_tree.py index 1d16083c..6f25ecb7 100644 --- a/indica/writers/bda_tree.py +++ b/indica/writers/bda_tree.py @@ -10,35 +10,29 @@ def bda(): nodes = { "TIME": ("NUMERIC", "time vector, s"), "TIME_OPT": ("NUMERIC", "time of optimisation, s"), - "INPUT": { "BURN_FRAC": ("NUMERIC", "Burn in fraction for chains"), "ITER": ("NUMERIC", "Iterations of optimiser"), "PARAM_NAMES": ("TEXT", "Names of parameters optimised"), - "OPT_QUANTITY":("TEXT", "Names of quantities optimised"), + "OPT_QUANTITY": ("TEXT", "Names of quantities optimised"), "MODEL_KWARGS": ("TEXT", "Model key word arguments"), # "OPT_KWARGS": ("TEXT", "Optimiser key word arguments"), "PULSE": ("NUMERIC", "Pulse number"), - "TSTART": ("NUMERIC", "Start of time vector, s"), "TEND": ("NUMERIC", "End of time vector, s"), "DT": ("NUMERIC", "Distance between time points, s"), "TSAMPLE": ("NUMERIC", "Sample time, s"), "IMPURITIES": ("TEXT", "Names of impurity elements"), "MAIN_ION": ("TEXT", "Name of main ion"), - }, - "METADATA": { "GITCOMMIT": ("TEXT", "Commit ID used for run"), "USER": ("TEXT", "Username of owner"), "EQUIL": ("TEXT", "Equilibrium used"), }, - "PROFILES": { "RHOP": ("NUMERIC", "Radial vector, Sqrt of normalised poloidal flux"), "RHOT": ("SIGNAL", "Radial vector, toroidal flux"), - "NE": ("SIGNAL", "Electron density, m^-3"), "NI": ("SIGNAL", "Ion density, m^-3"), "TE": ("SIGNAL", "Electron temperature, eV"), @@ -51,7 +45,6 @@ def bda(): "NIZ3": ("SIGNAL", "Density of impurity Z3, m^-3"), "NNEUTR": ("SIGNAL", "Density of neutral main ion, m^-3"), "NFAST": ("SIGNAL", "Density of fast ion, m^-3"), - "ZI": ("SIGNAL", "Average charge of main ion, "), "ZIM1": ("SIGNAL", "Average charge of impurity IMP1, "), "ZIM2": ("SIGNAL", "Average charge of impurity IMP2, "), @@ -59,7 +52,6 @@ def bda(): "ZEFF": ("SIGNAL", "Effective charge, "), "P": ("SIGNAL", "Pressure,Pa"), "VOLUME": ("SIGNAL", "Volume inside magnetic surface,m^3"), - "NE_ERR": ("SIGNAL", "Electron density error, m^-3"), "NI_ERR": ("SIGNAL", "Ion density error, m^-3"), "TE_ERR": ("SIGNAL", "Electron temperature error, eV"), @@ -72,9 +64,7 @@ def bda(): "NIZ3_ERR": ("SIGNAL", "Density of impurity Z3 error, m^-3"), "NNEUTR_ERR": ("SIGNAL", "Density of neutral main ion error, m^-3"), "NFAST_ERR": ("SIGNAL", "Density of fast ion error, m^-3"), - }, - "PROFILE_STAT": { "SAMPLES": ("NUMERIC", "Numerical index of the optimisation samples"), "RHOP": ("NUMERIC", "Radial vector, Sqrt of normalised poloidal flux"), @@ -90,7 +80,6 @@ def bda(): "NIZ3": ("SIGNAL", "Density of impurity IMP3, m^-3"), "NNEUTR": ("SIGNAL", "Density of neutral main ion, m^-3"), "NFAST": ("SIGNAL", "Density of fast ions, m^-3"), - "ZI": ("SIGNAL", "Average charge of main ion, "), "ZIM1": ("SIGNAL", "Average charge of impurity IMP1, "), "ZIM2": ("SIGNAL", "Average charge of impurity IMP2, "), @@ -99,7 +88,6 @@ def bda(): "P": ("SIGNAL", "Pressure,Pa"), "VOLUME": ("SIGNAL", "Volume inside magnetic surface, m^3"), }, - "GLOBAL": { "NE0": ("SIGNAL", "Central electron density, m^-3"), "NI0": ("SIGNAL", "Central ion density, m^-3"), @@ -111,7 +99,6 @@ def bda(): "NI0Z1": ("SIGNAL", "Central density of impurity Z1, m^-3"), "NI0Z2": ("SIGNAL", "Central density of impurity Z2, m^-3"), "NI0Z3": ("SIGNAL", "Central density of impurity Z3, m^-3"), - "NE0_ERR": ("SIGNAL", "Central electron density error, m^-3"), "NI0_ERR": ("SIGNAL", "Central ion density error, m^-3"), "TE0_ERR": ("SIGNAL", "Central electron temperature error, eV"), @@ -122,11 +109,13 @@ def bda(): "NI0Z1_ERR": ("SIGNAL", "Central density of impurity Z1, m^-3"), "NI0Z2_ERR": ("SIGNAL", "Central density of impurity Z2, m^-3"), "NI0Z3_ERR": ("SIGNAL", "Central density of impurity Z3, m^-3"), - }, "PHANTOMS": { "FLAG": ("TEXT", "True if phantoms used"), - "RHO_POLOIDAL": ("NUMERIC", "Radial vector, Sqrt of normalised poloidal flux"), + "RHO_POLOIDAL": ( + "NUMERIC", + "Radial vector, Sqrt of normalised poloidal flux", + ), "NE": ("SIGNAL", "Electron density, m^-3"), "NI": ("SIGNAL", "Ion density, m^-3"), "TE": ("SIGNAL", "Electron temperature, eV"), @@ -137,9 +126,7 @@ def bda(): "NIZ1": ("SIGNAL", "Impurity density of Z1, m^-3"), "NIZ2": ("SIGNAL", "Impurity density of Z2, m^-3"), "NIZ3": ("SIGNAL", "Impurity density of Z3, m^-3"), - }, - "OPTIMISATION": { "ACCEPT_FRAC": ("NUMERIC", "Fraction of samples accepted by optimiser"), "AUTO_CORR": ("NUMERIC", "Auto-correlation (iteration nwalker)"), @@ -150,37 +137,53 @@ def bda(): return nodes -DIAGNOSTIC_QUANTITY = ["DIAGNOSTIC1.QUANTITY1", "DIAGNOSTIC1.QUANTITY2", "DIAGNOSTIC2.QUANTITY1"] +DIAGNOSTIC_QUANTITY = [ + "DIAGNOSTIC1.QUANTITY1", + "DIAGNOSTIC1.QUANTITY2", + "DIAGNOSTIC2.QUANTITY1", +] -def create_nodes(pulse_to_write=23000101, run="RUN01", best=True, diagnostic_quantities=DIAGNOSTIC_QUANTITY, mode="EDIT"): +def create_nodes( + pulse_to_write=23000101, + run="RUN01", + best=True, + diagnostic_quantities=DIAGNOSTIC_QUANTITY, + mode="EDIT", +): bda_nodes = bda() - quant_list = [item.upper().split(".") for item in diagnostic_quantities] # replace OPTIMISED_QUANTITY + quant_list = [ + item.upper().split(".") for item in diagnostic_quantities + ] # replace OPTIMISED_QUANTITY diag_names = list(set([item[0] for item in quant_list])) - diag_nodes = {diag_name: - {quantity[1]: ("SIGNAL", f"measured {quantity[1]} from {quantity[0]}") - for quantity in quant_list if quantity[0] == diag_name} - for diag_name in diag_names - } + diag_nodes = { + diag_name: { + quantity[1]: ("SIGNAL", f"measured {quantity[1]} from {quantity[0]}") + for quantity in quant_list + if quantity[0] == diag_name + } + for diag_name in diag_names + } # for diag_name in diag_names: # diag_nodes[diag_name]["RUN"] = ("TEXT", f"RUN from which {diag_name} data was taken") nodes = { "RUN": ("TEXT", "RUN used for diagnostic"), "USAGE": ("TEXT", "Quantity used in analysis"), - "PULSE": ("NUMERIC", "Pulse used for diagnostic") + "PULSE": ("NUMERIC", "Pulse used for diagnostic"), } - workflow_nodes = {diag_name: nodes - for diag_name in diag_names - } + workflow_nodes = {diag_name: nodes for diag_name in diag_names} - model_nodes = {diag_name: - {quantity[1]: ("SIGNAL", f"modelled {quantity[1]} from {quantity[0]}") - for quantity in quant_list if quantity[0] == diag_name} - for diag_name in diag_names - } + model_nodes = { + diag_name: { + quantity[1]: ("SIGNAL", f"modelled {quantity[1]} from {quantity[0]}") + for quantity in quant_list + if quantity[0] == diag_name + } + for diag_name in diag_names + } model_nodes["SAMPLES"] = ("NUMERIC", "index of samples taken from optimisation") bda_nodes["MODEL_DATA"] = model_nodes bda_nodes["DIAG_DATA"] = diag_nodes @@ -201,7 +204,7 @@ def create_nodes(pulse_to_write=23000101, run="RUN01", best=True, diagnostic_qua tree=tree, pulse_number=pulse_to_write, base_node_to_read=None, - node_information_file=bda_nodes + node_information_file=bda_nodes, ).get() util.StandardNodeCreation( @@ -224,26 +227,33 @@ def write_nodes(pulse, node_info, data): ) -def check_analysis_run(pulseNo, which_run, ): +def check_analysis_run( + pulseNo, + which_run, +): # Checker function to see if data already exists in a run - IP_address_smaug = '192.168.1.7:8000' + IP_address_smaug = "192.168.1.7:8000" conn = Connection(IP_address_smaug) - conn.openTree('BDA', pulseNo) + conn.openTree("BDA", pulseNo) - temp = conn.get('BDA.' + which_run + '.TIME').data() + temp = conn.get("BDA." + which_run + ".TIME").data() conn.closeAllTrees() overwrite_flag = True if isinstance(temp, np.ndarray): - print(f'Data already Exists in pulseNo = {pulseNo}, which_run = {which_run}') - print('User prompt...') - question = f' Scheduled to overwrite pulseNo {pulseNo}, {which_run}' \ - f'\n Do you want to overwrite {which_run}? (y/n)' + print(f"Data already Exists in pulseNo = {pulseNo}, which_run = {which_run}") + print("User prompt...") + question = ( + f" Scheduled to overwrite pulseNo {pulseNo}, {which_run}" + f"\n Do you want to overwrite {which_run}? (y/n)" + ) overwrite_flag = query_yes_no(question) return overwrite_flag -def query_yes_no(question, ): +def query_yes_no( + question, +): valid = {"yes": True, "y": True, "no": False, "n": False} while True: sys.stdout.write(question) diff --git a/tests/unit/workflows/test_bayes_workflow_example.py b/tests/unit/workflows/test_bayes_workflow_example.py index e4033134..2da073e9 100644 --- a/tests/unit/workflows/test_bayes_workflow_example.py +++ b/tests/unit/workflows/test_bayes_workflow_example.py @@ -2,8 +2,13 @@ import copy import numpy as np -from indica.workflows.bayes_workflow_example import BayesWorkflowExample, DEFAULT_PRIORS, DEFAULT_PROFILE_PARAMS, \ - OPTIMISED_PARAMS, OPTIMISED_QUANTITY +from indica.workflows.bayes_workflow_example import ( + BayesWorkflowExample, + DEFAULT_PRIORS, + DEFAULT_PROFILE_PARAMS, + OPTIMISED_PARAMS, + OPTIMISED_QUANTITY, +) """ TODO: @@ -16,27 +21,28 @@ class TestBayesWorkflowExample: def setup_class(self): - - self.default_settings = dict(pulse=None, - pulse_to_write=23000111, - run="RUN01", - diagnostics=["xrcs", "efit", "smmh1", "cxff_pi"], - opt_quantity=OPTIMISED_QUANTITY, - param_names=OPTIMISED_PARAMS, - profile_params=DEFAULT_PROFILE_PARAMS, - priors=DEFAULT_PRIORS, - model_kwargs={"xrcs_moment_analysis": False, }, - phantoms=True, - - iterations=1, - nwalkers=20, - burn_frac=0.10, - dt=0.005, - tsample=0.060, - - mds_write=False, - plot=False, - sample_high_density=False, ) + self.default_settings = dict( + pulse=None, + pulse_to_write=23000111, + run="RUN01", + diagnostics=["xrcs", "efit", "smmh1", "cxff_pi"], + opt_quantity=OPTIMISED_QUANTITY, + param_names=OPTIMISED_PARAMS, + profile_params=DEFAULT_PROFILE_PARAMS, + priors=DEFAULT_PRIORS, + model_kwargs={ + "xrcs_moment_analysis": False, + }, + phantoms=True, + iterations=1, + nwalkers=20, + burn_frac=0.10, + dt=0.005, + tsample=0.060, + mds_write=False, + plot=False, + sample_high_density=False, + ) self.untouched_workflow = BayesWorkflowExample(**self.default_settings) self.workflow = copy.deepcopy(self.untouched_workflow) @@ -48,7 +54,15 @@ def teardown_method(self): self.workflow = None def test_workflow_initializes(self): - attributes_to_check = ["plasma", "phantom_profiles", "data", "reader", "opt_data", "models", "bayesmodel",] + attributes_to_check = [ + "plasma", + "phantom_profiles", + "data", + "reader", + "opt_data", + "models", + "bayesmodel", + ] for attribute in attributes_to_check: if not hasattr(self.workflow, attribute): raise ValueError(f"missing {attribute} in workflow object") @@ -56,11 +70,15 @@ def test_workflow_initializes(self): def test_init_phantoms_false_with_example_plasma(self): with pytest.raises(ValueError): - example = BayesWorkflowExample(dict(self.default_settings, **{"phantoms": False})) + example = BayesWorkflowExample( + dict(self.default_settings, **{"phantoms": False}) + ) def test_init_not_including_all_required_inputs(self): with pytest.raises(ValueError): - example = BayesWorkflowExample(dict(self.default_settings, **{"param_names": None})) + example = BayesWorkflowExample( + dict(self.default_settings, **{"param_names": None}) + ) # def test_reader_has_read_all_diagnostic_data(self): # assert all(diag_name in self.workflow.reader.keys() for diag_name in self.workflow.diagnostics) @@ -71,7 +89,7 @@ def test_plasma_has_equilibrium(self): def test_phantom_profiles_are_not_mutatable(self): phantoms = copy.deepcopy(self.workflow.phantom_profiles) self.workflow.plasma.electron_temperature += 1 - assert (phantoms is not self.workflow.phantom_profiles) + assert phantoms is not self.workflow.phantom_profiles def test_setup_models_with_wrong_diagnostic_names(self): with pytest.raises(Exception): @@ -113,7 +131,7 @@ def test_worklow_has_results_after_run(self): raise ValueError(f"missing result in workflow object") assert True + if __name__ == "__main__": test = TestBayesWorkflowExample() test.setup_class() - From acc50b574fbf3337f2b9c59266705354e70761d6 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 10 Aug 2023 10:44:00 +0100 Subject: [PATCH 120/160] renamed kin_prof -> plasma_profiles --- indica/bayesmodels.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/indica/bayesmodels.py b/indica/bayesmodels.py index c8801664..dd9e5410 100644 --- a/indica/bayesmodels.py +++ b/indica/bayesmodels.py @@ -70,7 +70,6 @@ def __init__( def _build_bckc(self, params, **kwargs): """ - TODO: consider how to handle when models have overlapping kwargs Parameters ---------- params - dictionary which is updated by optimiser @@ -82,14 +81,14 @@ def _build_bckc(self, params, **kwargs): """ self.bckc: dict = {} for model in self.diagnostic_models: - # removes "model.name_" from params and kwargs then passes them to model e.g. xrcs_background -> background + # removes "model.name." from params and kwargs then passes them to model e.g. xrcs.background -> background _nuisance_params = { - param_name.replace(model.name + "_", ""): param_value + param_name.replace(model.name + ".", ""): param_value for param_name, param_value in params.items() if model.name in param_name } _model_settings = { - kwarg_name.replace(model.name + "_", ""): kwarg_value + kwarg_name.replace(model.name + ".", ""): kwarg_value for kwarg_name, kwarg_value in kwargs.items() if model.name in kwarg_name } @@ -248,14 +247,14 @@ def ln_posterior(self, parameters: dict, **kwargs): ln_likelihood = self._ln_likelihood() # compare results to data ln_posterior = ln_likelihood + ln_prior - kin_profs = {} + plasma_profiles = {} for profile_key in PROFILES: if hasattr(self.plasma, profile_key): - kin_profs[profile_key] = getattr(self.plasma, profile_key).sel( + plasma_profiles[profile_key] = getattr(self.plasma, profile_key).sel( t=self.plasma.time_to_calculate ) else: raise ValueError(f"plasma does not have attribute {profile_key}") - blob = deepcopy({**self.bckc, **kin_profs}) + blob = deepcopy({**self.bckc, **plasma_profiles}) return ln_posterior, blob From 223a526f67a3339934a3bc49de6c4e6af2e633bc Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 10 Aug 2023 10:47:57 +0100 Subject: [PATCH 121/160] refactored window handling --- indica/models/helike_spectroscopy.py | 21 ++++++++++----------- 1 file changed, 10 insertions(+), 11 deletions(-) diff --git a/indica/models/helike_spectroscopy.py b/indica/models/helike_spectroscopy.py index 12642af3..db5c4d39 100644 --- a/indica/models/helike_spectroscopy.py +++ b/indica/models/helike_spectroscopy.py @@ -46,7 +46,7 @@ def __init__( element: str = "ar", window_len: int = 1030, window_lim: list = [0.394, 0.401], - window_vector=None, + window: np.array = None, window_masks: list = [], ): """ @@ -77,6 +77,8 @@ def __init__( self.full_run = full_run self.adf15 = ADF15 self.pecs: dict + self.window = window + self.window_masks = window_masks if self.marchuk: marchuck_reader = MARCHUKReader( @@ -87,16 +89,13 @@ def __init__( else: self.pecs = self._set_adas_pecs() - self.window_masks = window_masks - self.window_vector = window_vector - if self.window_vector is not None: - window = self.window_vector - else: - window = np.linspace(window_lim[0], window_lim[1], window_len) - mask = np.zeros(shape=window.shape) - if window_masks: - for mslice in window_masks: - mask[(window > mslice.start) & (window < mslice.stop)] = 1 + if self.window is None: + self.window = np.linspace(window_lim[0], window_lim[1], window_len) + + mask = np.zeros(shape=self.window.shape) + if self.window_masks: + for mslice in self.window_masks: + mask[(self.window > mslice.start) & (self.window < mslice.stop)] = 1 else: mask[:] = 1 From c9333ae75893d1ec581d52b1e4ad9a86eab94316 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 10 Aug 2023 15:11:39 +0100 Subject: [PATCH 122/160] replaced doppler_broaden with physics.ev_doppler --- indica/models/helike_spectroscopy.py | 26 ++++++-------------------- 1 file changed, 6 insertions(+), 20 deletions(-) diff --git a/indica/models/helike_spectroscopy.py b/indica/models/helike_spectroscopy.py index db5c4d39..083ca4e9 100644 --- a/indica/models/helike_spectroscopy.py +++ b/indica/models/helike_spectroscopy.py @@ -3,10 +3,10 @@ import matplotlib.cm as cm import matplotlib.pylab as plt import numpy as np -from scipy import constants import xarray as xr from xarray import DataArray +from indica.physics import ev_doppler from indica.converters.line_of_sight import LineOfSightTransform from indica.datatypes import ELEMENTS from indica.models.abstractdiagnostic import DiagnosticModel @@ -90,12 +90,12 @@ def __init__( self.pecs = self._set_adas_pecs() if self.window is None: - self.window = np.linspace(window_lim[0], window_lim[1], window_len) + window = np.linspace(window_lim[0], window_lim[1], window_len) - mask = np.zeros(shape=self.window.shape) + mask = np.zeros(shape=window.shape) if self.window_masks: for mslice in self.window_masks: - mask[(self.window > mslice.start) & (self.window < mslice.stop)] = 1 + mask[(window > mslice.start) & (window < mslice.stop)] = 1 else: mask[:] = 1 @@ -492,13 +492,8 @@ def _make_spectra( """ element = self.pecs["element"] - _spectra = doppler_broaden( - self.window[self.window > 0].wavelength, - self.intensity, - self.intensity.wavelength, - self.ion_mass, - self.Ti.sel(element=element), - ) + _sigma = ev_doppler(self.Ti.sel(element=element), self.ion_mass) * self.window.wavelength + _spectra = gaussian(self.window[self.window > 0].wavelength, self.intensity, self.intensity.wavelength, _sigma) _spectra = _spectra.sum("line_name") # extend spectra to same coords as self.window.wavelength with NaNs # to maintain same shape as mds data @@ -697,18 +692,9 @@ def __call__( return self.bckc -# fmt: off -def doppler_broaden(x, integral, center, ion_mass, ion_temp): - mass = (ion_mass * constants.proton_mass * constants.c ** 2) - sigma = np.sqrt(constants.e / mass * ion_temp) * center - gaussian_broadened = gaussian(x, integral, center, sigma, ) - return gaussian_broadened - - def gaussian(x, integral, center, sigma): return (integral / (sigma * np.sqrt(2 * np.pi)) * np.exp(-((x - center) ** 2) / (2 * sigma ** 2))) -# fmt: on def select_transition(adf15_data, transition: str, wavelength: float): From 80bc391c3f3b09c395bbac4eb09942dbd98a88c8 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 10 Aug 2023 15:13:30 +0100 Subject: [PATCH 123/160] removed methods from __init__ --- indica/workflows/abstract_bayes_workflow.py | 8 +- indica/workflows/bayes_workflow_example.py | 148 ++++++++++---------- 2 files changed, 76 insertions(+), 80 deletions(-) diff --git a/indica/workflows/abstract_bayes_workflow.py b/indica/workflows/abstract_bayes_workflow.py index 86057daa..d9ff7184 100644 --- a/indica/workflows/abstract_bayes_workflow.py +++ b/indica/workflows/abstract_bayes_workflow.py @@ -13,24 +13,21 @@ class AbstractBayesWorkflow(ABC): @abstractmethod def __init__( self, + pulse=None, phantoms=None, diagnostics=None, param_names=None, opt_quantity=None, priors=None, ): + self.pulse = pulse self.phantoms = phantoms self.diagnostics = diagnostics self.param_names = param_names self.opt_quantity = opt_quantity self.priors = priors - self.read_data(self.diagnostics) - self.setup_plasma() - self.save_phantom_profiles() self.setup_models(self.diagnostics) - self.setup_opt_data(self.phantoms) - self.setup_optimiser() def read_test_data(self, diagnostic_transforms: dict): # Used with phantom data for purposes of tests @@ -59,6 +56,7 @@ def setup_plasma(self): """ self.plasma = None self.plasma.set_equilibrium(self.reader.equilibrium) + self.save_phantom_profiles() @abstractmethod def setup_models(self, diagnostics: list): diff --git a/indica/workflows/bayes_workflow_example.py b/indica/workflows/bayes_workflow_example.py index 02b20ec1..9df0efa3 100644 --- a/indica/workflows/bayes_workflow_example.py +++ b/indica/workflows/bayes_workflow_example.py @@ -105,51 +105,23 @@ class BayesWorkflowExample(AbstractBayesWorkflow): def __init__( - self, - pulse: int = None, - pulse_to_write: bool = None, - run: str = "RUN01", - diagnostics: list = None, - param_names: list = None, - opt_quantity: list = None, - priors: dict = None, - profile_params: dict = None, - phantoms: bool = False, - model_kwargs: dict = {}, - nwalkers: int = 50, - tstart: float = 0.02, - tend: float = 0.10, - dt: float = 0.01, - tsample: float = 0.06, - iterations: int = 100, - burn_frac: float = 0, - mds_write: bool = False, - plot: bool = True, - sample_high_density: bool = False, + self, + pulse: int = None, + diagnostics: list = None, + param_names: list = None, + opt_quantity: list = None, + priors: dict = None, + profile_params: dict = None, + phantoms: bool = False, ): self.pulse = pulse - self.pulse_to_write = pulse_to_write - self.run = run self.diagnostics = diagnostics self.param_names = param_names self.opt_quantity = opt_quantity self.priors = priors self.profile_params = profile_params - self.model_kwargs = model_kwargs self.phantoms = phantoms - self.tstart = tstart - self.tend = tend - self.dt = dt - self.tsample = tsample - self.nwalkers = nwalkers - self.iterations = iterations - self.burn_frac = burn_frac - - self.mds_write = mds_write - self.plot = plot - self.sample_high_density = sample_high_density - for attribute in [ "param_names", "opt_quantity", @@ -165,6 +137,7 @@ def __init__( "Set phantoms to True when running phantom plasma i.e. pulse=None" ) + # TODO: Add some abstraction here if pulse is None: print("Running in test mode") self.read_test_data( @@ -176,17 +149,19 @@ def __init__( ) else: self.read_data(self.diagnostics) - self.setup_plasma() - self.save_phantom_profiles() self.setup_models(self.diagnostics) - self.setup_opt_data(phantoms=self.phantoms) - self.setup_optimiser(self.model_kwargs) - def setup_plasma(self): + def setup_plasma(self, + tstart=0.02, + tend=0.10, + dt=0.005, + tsample=0.050, + ): + # TODO: move to plasma.py self.plasma = Plasma( - tstart=self.tstart, - tend=self.tend, - dt=self.dt, + tstart=tstart, + tend=tend, + dt=dt, main_ion="h", impurities=("ar", "c"), impurity_concentration=( @@ -197,11 +172,12 @@ def setup_plasma(self): n_rad=20, ) self.plasma.time_to_calculate = self.plasma.t[ - np.abs(self.tsample - self.plasma.t).argmin() + np.abs(tsample - self.plasma.t).argmin() ] self.plasma.set_equilibrium(self.equilibrium) self.plasma.update_profiles(self.profile_params) self.plasma.build_atomic_data(calc_power_loss=False) + self.save_phantom_profiles() def setup_models(self, diagnostics: list): self.models = {} @@ -230,17 +206,17 @@ def setup_models(self, diagnostics: list): elif diag == "xrcs": los_transform = self.transforms[diag] los_transform.set_equilibrium(self.plasma.equilibrium) - window_vector = None + window = None if hasattr(self, "data"): if diag in self.data.keys(): - window_vector = ( - self.data[diag]["spectra"].wavelength.values * 0.1 + window = ( + self.data[diag]["spectra"].wavelength.values * 0.1 ) model = HelikeSpectrometer( name="xrcs", window_masks=[slice(0.394, 0.396)], - window_vector=window_vector, + window=window, ) model.set_los_transform(los_transform) @@ -254,7 +230,7 @@ def setup_models(self, diagnostics: list): model.set_transect_transform(transform) else: raise ValueError(f"{diag} not found in setup_models") - model.plasma = self.plasma + # model.plasma = self.plasma TODO: add this to the models somewhere self.models[diag] = model def setup_opt_data(self, phantoms=False): @@ -268,28 +244,28 @@ def _phantom_data(self, noise=False, noise_factor=0.1): if "smmh1" in self.diagnostics: opt_data["smmh1.ne"] = ( self.models["smmh1"]() - .pop("ne") - .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) + .pop("ne") + .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) ) if "xrcs" in self.diagnostics: opt_data["xrcs.spectra"] = ( self.models["xrcs"]() - .pop("spectra") - .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) + .pop("spectra") + .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) ) opt_data["xrcs.spectra"]["error"] = np.sqrt(opt_data["xrcs.spectra"]) if "cxff_pi" in self.diagnostics: cxrs_data = ( self.models["cxff_pi"]() - .pop("ti") - .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) + .pop("ti") + .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) ) opt_data["cxff_pi.ti"] = cxrs_data.where(cxrs_data.channel == 2, drop=True) if "efit" in self.diagnostics: opt_data["efit.wp"] = ( self.models["efit"]() - .pop("wp") - .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) + .pop("wp") + .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) ) if noise: @@ -315,7 +291,7 @@ def _exp_data(self): opt_data = flatdict.FlatDict(self.data, ".") if "xrcs" in self.diagnostics: opt_data["xrcs.spectra"]["wavelength"] = ( - opt_data["xrcs.spectra"].wavelength * 0.1 + opt_data["xrcs.spectra"].wavelength * 0.1 ) background = opt_data["xrcs.spectra"].where( (opt_data["xrcs.spectra"].wavelength < 0.392) @@ -335,7 +311,19 @@ def _exp_data(self): ) return opt_data - def setup_optimiser(self, model_kwargs): + def setup_optimiser(self, + model_kwargs, + iterations=200, + nwalkers=50, + burn_frac=0.10, + sample_high_density=False, + ): + self.model_kwargs = model_kwargs + self.iterations = iterations + self.nwalkers = nwalkers + self.burn_frac = burn_frac + self.sample_high_density = sample_high_density + self.bayesmodel = BayesModels( plasma=self.plasma, data=self.opt_data, @@ -369,11 +357,13 @@ def _sample_start_points(self, sample_high_density: bool = True): ) return start_points - def __call__(self, filepath="./results/test/", **kwargs): - if self.mds_write: + def __call__(self, filepath="./results/test/", run=None, mds_write=False, pulse_to_write=None, plot=False, + **kwargs): + if mds_write: # check_analysis_run(self.pulse, self.run) self.node_structure = create_nodes( - pulse_to_write=self.pulse_to_write, + pulse_to_write=pulse_to_write, + run=run, diagnostic_quantities=self.opt_quantity, mode="NEW", ) @@ -381,11 +371,11 @@ def __call__(self, filepath="./results/test/", **kwargs): self.run_sampler() self.save_pickle(filepath=filepath) - if self.plot: # currently requires result with DataArrays + if plot: # currently requires result with DataArrays plot_bayes_result(self.result, filepath) self.result = self.dict_of_dataarray_to_numpy(self.result) - if self.mds_write: + if mds_write: write_nodes(self.pulse_to_write, self.node_structure, self.result) return self.result @@ -394,26 +384,34 @@ def __call__(self, filepath="./results/test/", **kwargs): if __name__ == "__main__": run = BayesWorkflowExample( pulse=None, - pulse_to_write=23000101, - run="RUN01", + phantoms=True, diagnostics=["xrcs", "efit", "smmh1", "cxff_pi"], opt_quantity=OPTIMISED_QUANTITY, param_names=OPTIMISED_PARAMS, profile_params=DEFAULT_PROFILE_PARAMS, priors=DEFAULT_PRIORS, - model_kwargs={ - "xrcs_moment_analysis": False, - }, - phantoms=True, + ) + + run.setup_plasma( + tstart=0.02, + tend=0.10, + dt=0.005, + tsample=0.060, + ) + run.setup_opt_data(phantoms=run.phantoms) + run.setup_optimiser( iterations=200, nwalkers=50, burn_frac=0.10, - dt=0.005, - tsample=0.060, - mds_write=True, - plot=True, sample_high_density=False, + model_kwargs={ + "xrcs_moment_analysis": False, + } ) results = run( filepath="./results/test/", + pulse_to_write=23000101, + run="RUN01", + mds_write=True, + plot=True, ) From 3f5c089626aa66d3bccbc3edfbc4c17434b4c9dc Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Fri, 11 Aug 2023 14:00:37 +0100 Subject: [PATCH 124/160] stashing --- indica/workflows/bayes_workflow_example.py | 5 +- .../workflows/test_bayes_workflow_example.py | 51 ++++++++++++------- 2 files changed, 36 insertions(+), 20 deletions(-) diff --git a/indica/workflows/bayes_workflow_example.py b/indica/workflows/bayes_workflow_example.py index 9df0efa3..79c34df7 100644 --- a/indica/workflows/bayes_workflow_example.py +++ b/indica/workflows/bayes_workflow_example.py @@ -230,7 +230,7 @@ def setup_models(self, diagnostics: list): model.set_transect_transform(transform) else: raise ValueError(f"{diag} not found in setup_models") - # model.plasma = self.plasma TODO: add this to the models somewhere + model.plasma = self.plasma #TODO: add this to the models somewhere self.models[diag] = model def setup_opt_data(self, phantoms=False): @@ -390,6 +390,9 @@ def __call__(self, filepath="./results/test/", run=None, mds_write=False, pulse_ param_names=OPTIMISED_PARAMS, profile_params=DEFAULT_PROFILE_PARAMS, priors=DEFAULT_PRIORS, + tstart=0.02, + tend=0.10, + dt=0.005, ) run.setup_plasma( diff --git a/tests/unit/workflows/test_bayes_workflow_example.py b/tests/unit/workflows/test_bayes_workflow_example.py index 2da073e9..5861c49a 100644 --- a/tests/unit/workflows/test_bayes_workflow_example.py +++ b/tests/unit/workflows/test_bayes_workflow_example.py @@ -21,30 +21,41 @@ class TestBayesWorkflowExample: def setup_class(self): - self.default_settings = dict( + self.init_settings = dict( pulse=None, - pulse_to_write=23000111, - run="RUN01", + phantoms = True, diagnostics=["xrcs", "efit", "smmh1", "cxff_pi"], opt_quantity=OPTIMISED_QUANTITY, param_names=OPTIMISED_PARAMS, profile_params=DEFAULT_PROFILE_PARAMS, priors=DEFAULT_PRIORS, - model_kwargs={ - "xrcs_moment_analysis": False, - }, - phantoms=True, - iterations=1, - nwalkers=20, - burn_frac=0.10, + tstart=0.02, + tend=0.10, dt=0.005, - tsample=0.060, - mds_write=False, - plot=False, - sample_high_density=False, ) + self.plasma_settings = dict( + + tsample=0.060,) + + + self.opt_settings = dict( + model_kwargs = { + "xrcs_moment_analysis": False, + }, + iterations = 1, + nwalkers = 20, + burn_frac = 0.10, + sample_high_density = False, + ) + + self.call_settings = dict( + filepath=None, + pulse_to_write=23000101, + run="RUN01", + mds_write = False, + plot = False,) - self.untouched_workflow = BayesWorkflowExample(**self.default_settings) + self.untouched_workflow = BayesWorkflowExample(**self.init_settings) self.workflow = copy.deepcopy(self.untouched_workflow) def setup_method(self): @@ -55,13 +66,10 @@ def teardown_method(self): def test_workflow_initializes(self): attributes_to_check = [ - "plasma", - "phantom_profiles", "data", "reader", - "opt_data", "models", - "bayesmodel", + "equilibrium" ] for attribute in attributes_to_check: if not hasattr(self.workflow, attribute): @@ -84,9 +92,11 @@ def test_init_not_including_all_required_inputs(self): # assert all(diag_name in self.workflow.reader.keys() for diag_name in self.workflow.diagnostics) def test_plasma_has_equilibrium(self): + self.workflow.setup_plasma() assert hasattr(self.workflow.plasma, "equilibrium") def test_phantom_profiles_are_not_mutatable(self): + self.workflow.setup_plasma() phantoms = copy.deepcopy(self.workflow.phantom_profiles) self.workflow.plasma.electron_temperature += 1 assert phantoms is not self.workflow.phantom_profiles @@ -126,6 +136,9 @@ def test_sampling_from_priors(self): assert True def test_worklow_has_results_after_run(self): + self.workflow.setup_plasma() + self.workflow.setup_opt_data(self.workflow.phantoms) + self.workflow.setup_optimiser() self.workflow.run_sampler() if not hasattr(self.workflow, "result"): raise ValueError(f"missing result in workflow object") From 766825a8e5ce315fc3db4742272ce5551f2ed2ee Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 15 Aug 2023 10:58:02 +0100 Subject: [PATCH 125/160] setup_plasma now takes kwargs --- indica/workflows/bayes_workflow_example.py | 62 ++++++++++++++-------- 1 file changed, 39 insertions(+), 23 deletions(-) diff --git a/indica/workflows/bayes_workflow_example.py b/indica/workflows/bayes_workflow_example.py index 79c34df7..0ce4f073 100644 --- a/indica/workflows/bayes_workflow_example.py +++ b/indica/workflows/bayes_workflow_example.py @@ -113,6 +113,9 @@ def __init__( priors: dict = None, profile_params: dict = None, phantoms: bool = False, + tstart=0.02, + tend=0.10, + dt=0.005, ): self.pulse = pulse self.diagnostics = diagnostics @@ -121,6 +124,9 @@ def __init__( self.priors = priors self.profile_params = profile_params self.phantoms = phantoms + self.tstart = tstart + self.tend = tend + self.dt = dt for attribute in [ "param_names", @@ -152,25 +158,34 @@ def __init__( self.setup_models(self.diagnostics) def setup_plasma(self, - tstart=0.02, - tend=0.10, - dt=0.005, + tstart=None, + tend=None, + dt=None, tsample=0.050, + main_ion = "h", + impurities = ("ar", "c"), + impurity_concentration = (0.001, 0.04), + n_rad=20, + **kwargs ): + + if not all([tstart, tend, dt]): + tstart = self.tstart + tend = self.tend + dt = self.dt + # TODO: move to plasma.py self.plasma = Plasma( tstart=tstart, tend=tend, dt=dt, - main_ion="h", - impurities=("ar", "c"), - impurity_concentration=( - 0.001, - 0.04, - ), + main_ion=main_ion, + impurities=impurities, + impurity_concentration=impurity_concentration, full_run=False, - n_rad=20, + n_rad=n_rad, ) + self.tsample = tsample self.plasma.time_to_calculate = self.plasma.t[ np.abs(tsample - self.plasma.t).argmin() ] @@ -199,13 +214,13 @@ def setup_models(self, diagnostics: list): # machine_dimensions=machine_dims, # passes=2, # ) - los_transform.set_equilibrium(self.plasma.equilibrium) + los_transform.set_equilibrium(self.equilibrium) model = Interferometry(name=diag) model.set_los_transform(los_transform) elif diag == "xrcs": los_transform = self.transforms[diag] - los_transform.set_equilibrium(self.plasma.equilibrium) + los_transform.set_equilibrium(self.equilibrium) window = None if hasattr(self, "data"): if diag in self.data.keys(): @@ -225,21 +240,25 @@ def setup_models(self, diagnostics: list): elif diag == "cxff_pi": transform = self.transforms[diag] - transform.set_equilibrium(self.plasma.equilibrium) + transform.set_equilibrium(self.equilibrium) model = ChargeExchange(name=diag, element="ar") model.set_transect_transform(transform) else: raise ValueError(f"{diag} not found in setup_models") - model.plasma = self.plasma #TODO: add this to the models somewhere self.models[diag] = model - def setup_opt_data(self, phantoms=False): + def setup_opt_data(self, phantoms=False, **kwargs): + if not hasattr(self, "plasma"): + raise ValueError("Missing plasma object required for setup_opt_data") + for model in self.models.values(): # Maybe refactor here... + model.plasma = self.plasma + if phantoms: - self.opt_data = self._phantom_data() + self.opt_data = self._phantom_data(**kwargs) else: - self.opt_data = self._exp_data() + self.opt_data = self._exp_data(**kwargs) - def _phantom_data(self, noise=False, noise_factor=0.1): + def _phantom_data(self, noise=False, noise_factor=0.1, **kwargs): opt_data = {} if "smmh1" in self.diagnostics: opt_data["smmh1.ne"] = ( @@ -269,7 +288,7 @@ def _phantom_data(self, noise=False, noise_factor=0.1): ) if noise: - opt_data["smmh1.ne"] = opt_data["smmh1.ne"] + self.opt_data[ + opt_data["smmh1.ne"] = opt_data["smmh1.ne"] + opt_data[ "smmh1.ne" ].max().values * np.random.normal(0, noise_factor, None) opt_data["xrcs.spectra"] = opt_data["xrcs.spectra"] + np.random.normal( @@ -287,7 +306,7 @@ def _phantom_data(self, noise=False, noise_factor=0.1): ].max().values * np.random.normal(0, noise_factor, None) return opt_data - def _exp_data(self): + def _exp_data(self, **kwargs): opt_data = flatdict.FlatDict(self.data, ".") if "xrcs" in self.diagnostics: opt_data["xrcs.spectra"]["wavelength"] = ( @@ -396,9 +415,6 @@ def __call__(self, filepath="./results/test/", run=None, mds_write=False, pulse_ ) run.setup_plasma( - tstart=0.02, - tend=0.10, - dt=0.005, tsample=0.060, ) run.setup_opt_data(phantoms=run.phantoms) From cb93ae8c724cb14349de25aaa9734e7df5ec8cdc Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 15 Aug 2023 10:59:05 +0100 Subject: [PATCH 126/160] adding plasma to models now happens in setup_opt_data --- indica/workflows/abstract_bayes_workflow.py | 135 ++++++++++---------- 1 file changed, 69 insertions(+), 66 deletions(-) diff --git a/indica/workflows/abstract_bayes_workflow.py b/indica/workflows/abstract_bayes_workflow.py index d9ff7184..b40a9db8 100644 --- a/indica/workflows/abstract_bayes_workflow.py +++ b/indica/workflows/abstract_bayes_workflow.py @@ -12,13 +12,13 @@ class AbstractBayesWorkflow(ABC): @abstractmethod def __init__( - self, - pulse=None, - phantoms=None, - diagnostics=None, - param_names=None, - opt_quantity=None, - priors=None, + self, + pulse=None, + phantoms=None, + diagnostics=None, + param_names=None, + opt_quantity=None, + priors=None, ): self.pulse = pulse self.phantoms = phantoms @@ -81,6 +81,9 @@ def setup_opt_data(self, phantom: bool = False): """ Get and prepare the data in necessary format for optimiser """ + for model in self.models: + model.plasma = self.plasma + if phantom: self.opt_data = self._phantom_data() else: @@ -107,8 +110,8 @@ def save_phantom_profiles(self, kinetic_profiles=None): if self.phantoms: phantom_profiles = { profile_key: getattr(self.plasma, profile_key) - .sel(t=self.plasma.time_to_calculate) - .copy() + .sel(t=self.plasma.time_to_calculate) + .copy() for profile_key in kinetic_profiles } else: @@ -116,7 +119,7 @@ def save_phantom_profiles(self, kinetic_profiles=None): profile_key: getattr(self.plasma, profile_key).sel( t=self.plasma.time_to_calculate ) - * 0 + * 0 for profile_key in kinetic_profiles } @@ -201,7 +204,7 @@ def _build_result_dict(self): } result["PHANTOMS"].update( { - f"NIZ{num_imp+1}": self.phantom_profiles["impurity_density"].sel( + f"NIZ{num_imp + 1}": self.phantom_profiles["impurity_density"].sel( element=imp ) for num_imp, imp in enumerate(self.plasma.impurities) @@ -221,40 +224,40 @@ def _build_result_dict(self): "RHO_TOR": self.plasma.equilibrium.rhotor.interp(t=self.plasma.t), "NE": self.blobs["electron_density"].median(dim="index"), "NI": self.blobs["ion_density"] - .sel(element=self.plasma.main_ion) - .median(dim="index"), + .sel(element=self.plasma.main_ion) + .median(dim="index"), "TE": self.blobs["electron_temperature"].median(dim="index"), "TI": self.blobs["ion_temperature"] - .sel(element=self.plasma.main_ion) - .median(dim="index"), + .sel(element=self.plasma.main_ion) + .median(dim="index"), "NFAST": self.blobs["fast_density"].median(dim="index"), "NNEUTR": self.blobs["neutral_density"].median(dim="index"), "NE_ERR": self.blobs["electron_density"].std(dim="index"), "NI_ERR": self.blobs["ion_density"] - .sel(element=self.plasma.main_ion) - .std(dim="index"), + .sel(element=self.plasma.main_ion) + .std(dim="index"), "TE_ERR": self.blobs["electron_temperature"].std(dim="index"), "TI_ERR": self.blobs["ion_temperature"] - .sel(element=self.plasma.main_ion) - .std(dim="index"), + .sel(element=self.plasma.main_ion) + .std(dim="index"), "NFAST_ERR": self.blobs["fast_density"].std(dim="index"), "NNEUTR_ERR": self.blobs["neutral_density"].std(dim="index"), } result["PROFILES"] = { **result["PROFILES"], **{ - f"NIZ{num_imp+1}": self.blobs["impurity_density"] - .sel(element=imp) - .median(dim="index") + f"NIZ{num_imp + 1}": self.blobs["impurity_density"] + .sel(element=imp) + .median(dim="index") for num_imp, imp in enumerate(self.plasma.impurities) }, } result["PROFILES"] = { **result["PROFILES"], **{ - f"NIZ{num_imp+1}_ERR": self.blobs["impurity_density"] - .sel(element=imp) - .std(dim="index") + f"NIZ{num_imp + 1}_ERR": self.blobs["impurity_density"] + .sel(element=imp) + .std(dim="index") for num_imp, imp in enumerate(self.plasma.impurities) }, } @@ -262,8 +265,8 @@ def _build_result_dict(self): **result["PROFILES"], **{ f"TIZ{num_imp + 1}": self.blobs["ion_temperature"] - .sel(element=imp) - .median(dim="index") + .sel(element=imp) + .median(dim="index") for num_imp, imp in enumerate(self.plasma.impurities) }, } @@ -271,8 +274,8 @@ def _build_result_dict(self): **result["PROFILES"], **{ f"TIZ{num_imp + 1}_ERR": self.blobs["ion_temperature"] - .sel(element=imp) - .std(dim="index") + .sel(element=imp) + .std(dim="index") for num_imp, imp in enumerate(self.plasma.impurities) }, } @@ -290,7 +293,7 @@ def _build_result_dict(self): result["PROFILE_STAT"] = { **result["PROFILE_STAT"], **{ - f"NIZ{num_imp+1}": self.blobs["impurity_density"].sel(element=imp) + f"NIZ{num_imp + 1}": self.blobs["impurity_density"].sel(element=imp) for num_imp, imp in enumerate(self.plasma.impurities) }, } @@ -311,41 +314,41 @@ def _build_result_dict(self): result["GLOBAL"] = { "TI0": self.blobs["ion_temperature"] - .sel(element=self.plasma.main_ion) - .sel(rho_poloidal=0, method="nearest") - .median(dim="index"), + .sel(element=self.plasma.main_ion) + .sel(rho_poloidal=0, method="nearest") + .median(dim="index"), "TE0": self.blobs["electron_temperature"] - .sel(rho_poloidal=0, method="nearest") - .median(dim="index"), + .sel(rho_poloidal=0, method="nearest") + .median(dim="index"), "NE0": self.blobs["electron_density"] - .sel(rho_poloidal=0, method="nearest") - .median(dim="index"), + .sel(rho_poloidal=0, method="nearest") + .median(dim="index"), "NI0": self.blobs["ion_density"] - .sel(element=self.plasma.main_ion) - .sel(rho_poloidal=0, method="nearest") - .median(dim="index"), + .sel(element=self.plasma.main_ion) + .sel(rho_poloidal=0, method="nearest") + .median(dim="index"), "TI0_ERR": self.blobs["ion_temperature"] - .sel(element=self.plasma.main_ion) - .sel(rho_poloidal=0, method="nearest") - .std(dim="index"), + .sel(element=self.plasma.main_ion) + .sel(rho_poloidal=0, method="nearest") + .std(dim="index"), "TE0_ERR": self.blobs["electron_temperature"] - .sel(rho_poloidal=0, method="nearest") - .std(dim="index"), + .sel(rho_poloidal=0, method="nearest") + .std(dim="index"), "NE0_ERR": self.blobs["electron_density"] - .sel(rho_poloidal=0, method="nearest") - .std(dim="index"), + .sel(rho_poloidal=0, method="nearest") + .std(dim="index"), "NI0_ERR": self.blobs["ion_density"] - .sel(element=self.plasma.main_ion) - .sel(rho_poloidal=0, method="nearest") - .std(dim="index"), + .sel(element=self.plasma.main_ion) + .sel(rho_poloidal=0, method="nearest") + .std(dim="index"), } result["GLOBAL"] = { **result["GLOBAL"], **{ f"TI0Z{num_imp + 1}": self.blobs["ion_temperature"] - .sel(element=imp) - .sel(rho_poloidal=0, method="nearest") - .median(dim="index") + .sel(element=imp) + .sel(rho_poloidal=0, method="nearest") + .median(dim="index") for num_imp, imp in enumerate(self.plasma.impurities) }, } @@ -353,9 +356,9 @@ def _build_result_dict(self): **result["GLOBAL"], **{ f"TI0Z{num_imp + 1}_ERR": self.blobs["ion_temperature"] - .sel(element=imp) - .sel(rho_poloidal=0, method="nearest") - .std(dim="index") + .sel(element=imp) + .sel(rho_poloidal=0, method="nearest") + .std(dim="index") for num_imp, imp in enumerate(self.plasma.impurities) }, } @@ -363,9 +366,9 @@ def _build_result_dict(self): **result["GLOBAL"], **{ f"NI0Z{num_imp + 1}": self.blobs["impurity_density"] - .sel(element=imp) - .sel(rho_poloidal=0, method="nearest") - .median(dim="index") + .sel(element=imp) + .sel(rho_poloidal=0, method="nearest") + .median(dim="index") for num_imp, imp in enumerate(self.plasma.impurities) }, } @@ -373,9 +376,9 @@ def _build_result_dict(self): **result["GLOBAL"], **{ f"NI0Z{num_imp + 1}_ERR": self.blobs["impurity_density"] - .sel(element=imp) - .sel(rho_poloidal=0, method="nearest") - .std(dim="index") + .sel(element=imp) + .sel(rho_poloidal=0, method="nearest") + .std(dim="index") for num_imp, imp in enumerate(self.plasma.impurities) }, } @@ -452,9 +455,9 @@ def sample_with_autocorr(sampler, start_points, iterations, n_params, auto_sampl autocorr = np.ones(shape=(iterations, n_params)) * np.nan old_tau = np.inf for sample in sampler.sample( - start_points, - iterations=iterations, - progress=True, + start_points, + iterations=iterations, + progress=True, ): if sampler.iteration % auto_sample: continue @@ -465,5 +468,5 @@ def sample_with_autocorr(sampler, start_points, iterations, n_params, auto_sampl if converged: break old_tau = new_tau - autocorr = autocorr[: sampler.iteration,] + autocorr = autocorr[: sampler.iteration, ] return autocorr From cf1c127c6128b38200011131be7a56df09d82841 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 15 Aug 2023 10:59:28 +0100 Subject: [PATCH 127/160] workflow broken up into methods --- .../workflows/test_bayes_workflow_example.py | 66 +++++++++++-------- 1 file changed, 37 insertions(+), 29 deletions(-) diff --git a/tests/unit/workflows/test_bayes_workflow_example.py b/tests/unit/workflows/test_bayes_workflow_example.py index 5861c49a..b80fc392 100644 --- a/tests/unit/workflows/test_bayes_workflow_example.py +++ b/tests/unit/workflows/test_bayes_workflow_example.py @@ -12,10 +12,7 @@ """ TODO: - -test kwarg handling - -Mock reader for testing experimental data methods +Mock reader for testing experimental data reading """ @@ -34,18 +31,18 @@ def setup_class(self): dt=0.005, ) self.plasma_settings = dict( - - tsample=0.060,) + tsample=0.060, + ) - self.opt_settings = dict( + self.optimiser_settings = dict( model_kwargs = { "xrcs_moment_analysis": False, }, - iterations = 1, - nwalkers = 20, - burn_frac = 0.10, - sample_high_density = False, + iterations = 1, + nwalkers = 20, + burn_frac = 0.10, + sample_high_density = False, ) self.call_settings = dict( @@ -55,11 +52,11 @@ def setup_class(self): mds_write = False, plot = False,) - self.untouched_workflow = BayesWorkflowExample(**self.init_settings) - self.workflow = copy.deepcopy(self.untouched_workflow) + self.workflow_untouched = BayesWorkflowExample(**self.init_settings) + self.workflow=None def setup_method(self): - self.workflow = copy.deepcopy(self.untouched_workflow) + self.workflow = copy.deepcopy(self.workflow_untouched) def teardown_method(self): self.workflow = None @@ -79,34 +76,39 @@ def test_workflow_initializes(self): def test_init_phantoms_false_with_example_plasma(self): with pytest.raises(ValueError): example = BayesWorkflowExample( - dict(self.default_settings, **{"phantoms": False}) + dict(self.init_settings, **{"phantoms": False}) ) def test_init_not_including_all_required_inputs(self): with pytest.raises(ValueError): example = BayesWorkflowExample( - dict(self.default_settings, **{"param_names": None}) + dict(self.init_settings, **{"param_names": None}) ) # def test_reader_has_read_all_diagnostic_data(self): # assert all(diag_name in self.workflow.reader.keys() for diag_name in self.workflow.diagnostics) def test_plasma_has_equilibrium(self): - self.workflow.setup_plasma() + self.workflow.setup_plasma(**self.plasma_settings) assert hasattr(self.workflow.plasma, "equilibrium") def test_phantom_profiles_are_not_mutatable(self): - self.workflow.setup_plasma() + self.workflow.setup_plasma(**self.plasma_settings) phantoms = copy.deepcopy(self.workflow.phantom_profiles) self.workflow.plasma.electron_temperature += 1 assert phantoms is not self.workflow.phantom_profiles def test_setup_models_with_wrong_diagnostic_names(self): - with pytest.raises(Exception): + with pytest.raises(ValueError): self.workflow.setup_models(["foo", "bar", "xrcs"]) + def test_opt_data_without_plasma(self): + with pytest.raises(ValueError): + self.workflow.setup_opt_data(phantoms=True) + def test_phantom_data_exists(self): - self.workflow._phantom_data() + self.workflow.setup_plasma(**self.plasma_settings) + self.workflow.setup_opt_data(phantoms=True) assert self.workflow.opt_data # def test_experimental_data_exists(self): @@ -114,7 +116,8 @@ def test_phantom_data_exists(self): # assert self.workflow.opt_data def test_phantom_data_has_time_dim(self): - self.workflow._phantom_data() + self.workflow.setup_plasma(**self.plasma_settings) + self.workflow.setup_opt_data(phantoms=True) for key, value in self.workflow.opt_data.items(): assert "t" in value.dims @@ -124,21 +127,26 @@ def test_phantom_data_has_time_dim(self): # assert "t" in value.dims def test_phantom_data_runs_with_noise_added(self): - self.workflow._phantom_data(noise=True) + self.workflow.setup_plasma(**self.plasma_settings) + self.workflow.setup_opt_data(phantoms=True, noise=True) assert self.workflow.opt_data - def test_sampling_from_high_density(self): - self.workflow._sample_start_points(sample_high_density=True) + def test_sampling_from_priors(self): + self.workflow.setup_plasma(**self.plasma_settings) + self.workflow.setup_opt_data(phantoms=True, ) + self.workflow.setup_optimiser(**dict(self.optimiser_settings, **{"sample_high_density":False})) assert True - def test_sampling_from_priors(self): - self.workflow._sample_start_points(sample_high_density=False) + def test_sampling_from_high_density(self): + self.workflow.setup_plasma(**self.plasma_settings) + self.workflow.setup_opt_data(phantoms=True, ) + self.workflow.setup_optimiser(**dict(self.optimiser_settings, **{"sample_high_density":True})) assert True def test_worklow_has_results_after_run(self): - self.workflow.setup_plasma() - self.workflow.setup_opt_data(self.workflow.phantoms) - self.workflow.setup_optimiser() + self.workflow.setup_plasma(**self.plasma_settings) + self.workflow.setup_opt_data(phantoms=True) + self.workflow.setup_optimiser(**self.optimiser_settings) self.workflow.run_sampler() if not hasattr(self.workflow, "result"): raise ValueError(f"missing result in workflow object") From 260a5150812f4345246add54d20b217b332d4e58 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 15 Aug 2023 14:08:48 +0100 Subject: [PATCH 128/160] removed redundant kwargs --- indica/models/helike_spectroscopy.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/indica/models/helike_spectroscopy.py b/indica/models/helike_spectroscopy.py index b2dcb6fc..bfb7c406 100644 --- a/indica/models/helike_spectroscopy.py +++ b/indica/models/helike_spectroscopy.py @@ -4,7 +4,6 @@ import xarray as xr from xarray import DataArray -from indica.physics import ev_doppler from indica.converters.line_of_sight import LineOfSightTransform from indica.datatypes import ELEMENTS from indica.models.abstractdiagnostic import DiagnosticModel @@ -38,8 +37,6 @@ def __init__( instrument_method="get_helike_spectroscopy", etendue: float = 1.0, calibration: float = 8.0e-20, - marchuk: bool = True, - full_run: bool = False, element: str = "ar", window_len: int = 1030, window_lim: list = [0.394, 0.401], @@ -70,7 +67,8 @@ def __init__( self.line_ranges = LINE_RANGES self.line_labels = line_labels - window = np.linspace(window_lim[0], window_lim[1], window_len) + if window is None: + window = np.linspace(window_lim[0], window_lim[1], window_len) mask = np.zeros(shape=window.shape) if self.window_masks: for mslice in self.window_masks: From 0821b9dc62b70471493ae5cfa368e99dfb147b86 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 15 Aug 2023 14:40:50 +0100 Subject: [PATCH 129/160] fixed non_plasma call option --- indica/models/helike_spectroscopy.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/indica/models/helike_spectroscopy.py b/indica/models/helike_spectroscopy.py index bfb7c406..87ce54d2 100644 --- a/indica/models/helike_spectroscopy.py +++ b/indica/models/helike_spectroscopy.py @@ -39,10 +39,10 @@ def __init__( calibration: float = 8.0e-20, element: str = "ar", window_len: int = 1030, - window_lim: list = [0.394, 0.401], + window_lim=None, window: np.array = None, - window_masks: list = [], - line_labels: list = ["w", "k", "n3", "n345", "z", "qra"], + window_masks=None, + line_labels=None, ): """ Read all atomic data and initialise objects @@ -53,14 +53,19 @@ def __init__( String identifier for the spectrometer """ + if window_lim is None: + window_lim = [0.394, 0.401] + if window_masks is None: + window_masks = [] + if line_labels is None: + line_labels = ["w", "k", "n3", "n345", "z", "qra"] + self.name = name self.instrument_method = instrument_method - self.element: str = element z_elem, a_elem, name_elem = ELEMENTS[element] self.ion_charge: int = z_elem - 2 # He-like self.ion_mass: float = a_elem - self.etendue = etendue self.calibration = calibration self.window_masks = window_masks @@ -408,7 +413,7 @@ def __call__( "moment_analysis cannot be used when window_masks is not set to None" ) - if self.plasma is not None: + if hasattr(self, "plasma"): if t is None: t = self.plasma.time_to_calculate Te = self.plasma.electron_temperature.interp( From 78800c6704b0219fe427de419e89240bdf18d9ac Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 15 Aug 2023 15:22:00 +0100 Subject: [PATCH 130/160] fixed plasma_initialisation --- tests/unit/models/test_plasma.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/models/test_plasma.py b/tests/unit/models/test_plasma.py index 9b2a8fb7..4746e702 100644 --- a/tests/unit/models/test_plasma.py +++ b/tests/unit/models/test_plasma.py @@ -35,7 +35,7 @@ def setup_class(self): # return def teardown_method(self): - self.plasma.initialize_variables(tstart=self.tstart, tend=self.tend, dt=self.dt) + self.plasma.initialize_variables() def test_plasma_initializes(self): assert hasattr(self, "plasma") From 62d67277d4e6183c6e5ac407b8ae1b1374974d3c Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 15 Aug 2023 16:24:29 +0100 Subject: [PATCH 131/160] fixing violin plotting --- indica/workflows/bayes_plots.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/indica/workflows/bayes_plots.py b/indica/workflows/bayes_plots.py index 8695969d..c027c004 100644 --- a/indica/workflows/bayes_plots.py +++ b/indica/workflows/bayes_plots.py @@ -172,16 +172,18 @@ def violinplot( # quantiles=[0.025, 0.975, 0.16, 0.84], # showmedians=True, ) + y = diag_data[key].sel(t=data[key].t).values + # TODO abstract the yerr + axs.errorbar( + 1, y=y, yerr=y * 0.10, fmt="D", ecolor="black", capsize=10, color="black" + ) violin["bodies"][0].set_edgecolor("black") axs.set_xlabel(key) top = axs.get_ylim()[1] bot = axs.get_ylim()[0] axs.set_ylim(top=top * 1.1, bottom=bot * 0.9) axs.set_ylabel(f"{ylabel}") - y = diag_data[key].sel(t=data[key].t).values - axs.errorbar( - 1, y=y, yerr=y * 0.10, fmt="D", ecolor="black", capsize=10, color="black" - ) + set_axis_sci() plt.setp([axs.get_xticklabels()], visible=False) plt.savefig(figheader + filename) From 98190dda0476bb01c175cefc6ed56e1ce6e9a935 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Tue, 15 Aug 2023 16:44:34 +0100 Subject: [PATCH 132/160] units for xrcs.spectra.wavelength fixed --- indica/workflows/bayes_workflow_example.py | 48 +++++++++++----------- 1 file changed, 24 insertions(+), 24 deletions(-) diff --git a/indica/workflows/bayes_workflow_example.py b/indica/workflows/bayes_workflow_example.py index 0ce4f073..e2e12be5 100644 --- a/indica/workflows/bayes_workflow_example.py +++ b/indica/workflows/bayes_workflow_example.py @@ -7,6 +7,7 @@ from indica.workflows.bayes_plots import plot_bayes_result from indica.workflows.abstract_bayes_workflow import AbstractBayesWorkflow from indica.writers.bda_tree import create_nodes, write_nodes, check_analysis_run +from indica.converters.line_of_sight import LineOfSightTransform from indica.models.interferometry import Interferometry, smmh1_transform_example from indica.models.helike_spectroscopy import ( @@ -97,7 +98,7 @@ OPTIMISED_QUANTITY = [ "xrcs.spectra", - "cxff_pi.ti", + # "cxff_pi.ti", "efit.wp", "smmh1.ne", ] @@ -198,22 +199,22 @@ def setup_models(self, diagnostics: list): self.models = {} for diag in diagnostics: if diag == "smmh1": - los_transform = self.transforms[diag] - # machine_dims = self.plasma.machine_dimensions - # origin = np.array([[-0.38063365, 0.91893092, 0.01]]) - # # end = np.array([[0, 0, 0.01]]) - # direction = np.array([[0.38173721, -0.92387953, -0.02689453]]) - # los_transform = LineOfSightTransform( - # origin[:, 0], - # origin[:, 1], - # origin[:, 2], - # direction[:, 0], - # direction[:, 1], - # direction[:, 2], - # name="", - # machine_dimensions=machine_dims, - # passes=2, - # ) + # los_transform = self.transforms[diag] + machine_dims = ((0.15, 0.95), (-0.7, 0.7)) + origin = np.array([[-0.38063365, 0.91893092, 0.01]]) + # end = np.array([[0, 0, 0.01]]) + direction = np.array([[0.38173721, -0.92387953, -0.02689453]]) + los_transform = LineOfSightTransform( + origin[:, 0], + origin[:, 1], + origin[:, 2], + direction[:, 0], + direction[:, 1], + direction[:, 2], + name="", + machine_dimensions=machine_dims, + passes=2, + ) los_transform.set_equilibrium(self.equilibrium) model = Interferometry(name=diag) model.set_los_transform(los_transform) @@ -225,7 +226,7 @@ def setup_models(self, diagnostics: list): if hasattr(self, "data"): if diag in self.data.keys(): window = ( - self.data[diag]["spectra"].wavelength.values * 0.1 + self.data[diag]["spectra"].wavelength.values ) model = HelikeSpectrometer( @@ -395,16 +396,16 @@ def __call__(self, filepath="./results/test/", run=None, mds_write=False, pulse_ self.result = self.dict_of_dataarray_to_numpy(self.result) if mds_write: - write_nodes(self.pulse_to_write, self.node_structure, self.result) + write_nodes(pulse_to_write, self.node_structure, self.result) return self.result if __name__ == "__main__": run = BayesWorkflowExample( - pulse=None, + pulse=9780, phantoms=True, - diagnostics=["xrcs", "efit", "smmh1", "cxff_pi"], + diagnostics=["xrcs", "efit", "smmh1",], opt_quantity=OPTIMISED_QUANTITY, param_names=OPTIMISED_PARAMS, profile_params=DEFAULT_PROFILE_PARAMS, @@ -419,12 +420,11 @@ def __call__(self, filepath="./results/test/", run=None, mds_write=False, pulse_ ) run.setup_opt_data(phantoms=run.phantoms) run.setup_optimiser( - iterations=200, + iterations=2, nwalkers=50, burn_frac=0.10, - sample_high_density=False, + sample_high_density=True, model_kwargs={ - "xrcs_moment_analysis": False, } ) results = run( From f732da034d92817e50b9e5e239501d305113e355 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 17 Aug 2023 09:25:10 +0100 Subject: [PATCH 133/160] nchannels added to example_los --- indica/models/interferometry.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/indica/models/interferometry.py b/indica/models/interferometry.py index 2f6393e0..c9354205 100644 --- a/indica/models/interferometry.py +++ b/indica/models/interferometry.py @@ -89,9 +89,12 @@ def __call__( return self.bckc -def smmh1_transform_example(): - los_start = np.array([[0.8, 0, 0], [0.8, 0, -0.1], [0.8, 0, -0.2]]) - los_end = np.array([[0.17, 0, 0], [0.17, 0, -0.25], [0.17, 0, -0.2]]) +def smmh1_transform_example(nchannels): + + los_start = np.array([[0.8, 0, 0]]) * np.ones((nchannels,3)) + los_start[:,2] = np.linspace(0, -0.2, nchannels) + los_end = np.array([[0.17, 0, 0]]) * np.ones((nchannels, 3)) + los_end[:, 2] = np.linspace(0, -0.2, nchannels) origin = los_start direction = los_end - los_start los_transform = LineOfSightTransform( @@ -117,7 +120,7 @@ def example_run(pulse: int = None, plasma=None, plot=False): model = Interferometry( diagnostic_name, ) - los_transform = smmh1_transform_example() + los_transform = smmh1_transform_example(3) los_transform.set_equilibrium(plasma.equilibrium) model.set_los_transform(los_transform) model.set_plasma(plasma) From ec1c9cadd5d1b2231772239142d4b139c8267b3e Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 17 Aug 2023 09:26:41 +0100 Subject: [PATCH 134/160] adding print message for fake data reading --- indica/workflows/abstract_bayes_workflow.py | 1 + 1 file changed, 1 insertion(+) diff --git a/indica/workflows/abstract_bayes_workflow.py b/indica/workflows/abstract_bayes_workflow.py index b40a9db8..71b886ad 100644 --- a/indica/workflows/abstract_bayes_workflow.py +++ b/indica/workflows/abstract_bayes_workflow.py @@ -31,6 +31,7 @@ def __init__( def read_test_data(self, diagnostic_transforms: dict): # Used with phantom data for purposes of tests + print("Reading fake data") self.reader = None self.equilibrium = fake_equilibrium( self.tstart, From b0fac373937554d9938a5f0285f5cb7ec907fcac Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 17 Aug 2023 09:59:50 +0100 Subject: [PATCH 135/160] adjusting example los --- indica/models/helike_spectroscopy.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/indica/models/helike_spectroscopy.py b/indica/models/helike_spectroscopy.py index 87ce54d2..5fd32aae 100644 --- a/indica/models/helike_spectroscopy.py +++ b/indica/models/helike_spectroscopy.py @@ -465,12 +465,13 @@ def __call__( self._build_bckc_dictionary() return self.bckc -def helike_transform_example(nchannels=3): +def helike_transform_example(nchannels): los_end = np.full((nchannels, 3), 0.0) los_end[:, 0] = 0.17 los_end[:, 1] = 0.0 - los_end[:, 2] = np.linspace(0.43, -0.43, nchannels) - los_start = np.array([[0.8, 0, 0]] * los_end.shape[0]) + los_end[:, 2] = np.linspace(0.2, -0.5, nchannels) + los_start = np.array([[0.9, 0, 0]] * los_end.shape[0]) + los_start[:, 2] = -0.1 origin = los_start direction = los_end - los_start @@ -531,8 +532,7 @@ def example_run( plt.ylabel("spectra") plt.legend() - # model.los_transform.plot(tplot, - # plot_all=model.los_transform.x1.__len__() > 1) + los_transform.plot() if "int_w" in bckc.keys() & "t" in bckc["int_w"].dims: plt.figure() From 03d68c70efc37bc1a50e1b981827dfd6b08a7ebd Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 17 Aug 2023 10:12:45 +0100 Subject: [PATCH 136/160] bckc method not printing spectra/fit not available everytime --- indica/models/charge_exchange.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/indica/models/charge_exchange.py b/indica/models/charge_exchange.py index 9d223131..5551d8a6 100644 --- a/indica/models/charge_exchange.py +++ b/indica/models/charge_exchange.py @@ -42,6 +42,12 @@ def _build_bckc_dictionary(self): self.bckc[quantity] = self.Ti_at_channels long_name = "Ion temperature" units = "eV" + elif quant == "spectra": + # Placeholder + continue + elif quant == "fit": + # Placeholder + continue else: print(f"{quant} not available in model for {self.instrument_method}") continue From 72e6e7eb306a2545abedd634e290b2472aa7c8b5 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 17 Aug 2023 10:41:38 +0100 Subject: [PATCH 137/160] adjusting priors --- indica/workflows/bayes_workflow_example.py | 84 +++++++++++----------- 1 file changed, 43 insertions(+), 41 deletions(-) diff --git a/indica/workflows/bayes_workflow_example.py b/indica/workflows/bayes_workflow_example.py index e2e12be5..5d2e5868 100644 --- a/indica/workflows/bayes_workflow_example.py +++ b/indica/workflows/bayes_workflow_example.py @@ -26,8 +26,8 @@ "Ne_prof.y1": 2e18, "Ne_prof.yend": 1e18, "Ne_prof.wped": 2, - "Nimp_prof.y0": 3e16, - "Nimp_prof.y1": 0.5e16, + "Nimp_prof.y0": 1e17, + "Nimp_prof.y1": 5e15, "Nimp_prof.wcenter": 0.4, "Nimp_prof.wped": 6, "Nimp_prof.peaking": 2, @@ -50,28 +50,30 @@ "Ne_prof.wped": get_uniform(1, 6), "Ne_prof.wcenter": get_uniform(0.1, 0.8), "Ne_prof.peaking": get_uniform(1, 6), - "ar_conc": loguniform(0.0001, 0.01), + "Nimp_prof.y0": loguniform(1e16, 1e18), - "Nimp_prof.y1": get_uniform(1e15, 2e16), + "Nimp_prof.y1": get_uniform(1e15, 1e16), "Ne_prof.y0/Nimp_prof.y0": lambda x1, x2: np.where( (x1 > x2 * 100) & (x1 < x2 * 1e5), 1, 0 ), "Nimp_prof.y0/Nimp_prof.y1": lambda x1, x2: np.where((x1 > x2), 1, 0), "Nimp_prof.wped": get_uniform(1, 6), "Nimp_prof.wcenter": get_uniform(0.1, 0.8), - "Nimp_prof.peaking": get_uniform(1, 20), + "Nimp_prof.peaking": get_uniform(1, 10), "Nimp_prof.peaking/Ne_prof.peaking": lambda x1, x2: np.where( (x1 > x2), 1, 0 ), # impurity always more peaked + "Te_prof.y0": get_uniform(1000, 5000), "Te_prof.wped": get_uniform(1, 6), - "Te_prof.wcenter": get_uniform(0.1, 0.6), - "Te_prof.peaking": get_uniform(1, 6), - "Ti_prof.y0/Te_prof.y0": lambda x1, x2: np.where(x1 > x2, 1, 0), # hot ion mode - "Ti_prof.y0": get_uniform(3000, 10000), + "Te_prof.wcenter": get_uniform(0.1, 0.8), + "Te_prof.peaking": get_uniform(1, 10), + + # "Ti_prof.y0/Te_prof.y0": lambda x1, x2: np.where(x1 > x2, 1, 0), # hot ion mode + "Ti_prof.y0": get_uniform(1000, 10000), "Ti_prof.wped": get_uniform(1, 6), - "Ti_prof.wcenter": get_uniform(0.1, 0.6), - "Ti_prof.peaking": get_uniform(1, 20), + "Ti_prof.wcenter": get_uniform(0.1, 0.8), + "Ti_prof.peaking": get_uniform(1, 10), } OPTIMISED_PARAMS = [ @@ -80,29 +82,24 @@ # "Ne_prof.peaking", # "Ne_prof.wcenter", # "Ne_prof.wped", - # "ar_conc", + # "Nimp_prof.y1", "Nimp_prof.y0", # "Nimp_prof.wcenter", # "Nimp_prof.wped", # "Nimp_prof.peaking", + "Te_prof.y0", # "Te_prof.wped", "Te_prof.wcenter", "Te_prof.peaking", + "Ti_prof.y0", - "Ti_prof.wped", + # "Ti_prof.wped", "Ti_prof.wcenter", "Ti_prof.peaking", ] -OPTIMISED_QUANTITY = [ - "xrcs.spectra", - # "cxff_pi.ti", - "efit.wp", - "smmh1.ne", -] - class BayesWorkflowExample(AbstractBayesWorkflow): def __init__( @@ -150,7 +147,7 @@ def __init__( self.read_test_data( { "xrcs": helike_transform_example(1), - "smmh1": smmh1_transform_example(), + "smmh1": smmh1_transform_example(1), "cxff_pi": pi_transform_example(5), } ) @@ -199,22 +196,22 @@ def setup_models(self, diagnostics: list): self.models = {} for diag in diagnostics: if diag == "smmh1": - # los_transform = self.transforms[diag] - machine_dims = ((0.15, 0.95), (-0.7, 0.7)) - origin = np.array([[-0.38063365, 0.91893092, 0.01]]) - # end = np.array([[0, 0, 0.01]]) - direction = np.array([[0.38173721, -0.92387953, -0.02689453]]) - los_transform = LineOfSightTransform( - origin[:, 0], - origin[:, 1], - origin[:, 2], - direction[:, 0], - direction[:, 1], - direction[:, 2], - name="", - machine_dimensions=machine_dims, - passes=2, - ) + los_transform = self.transforms[diag] + # machine_dims = ((0.15, 0.95), (-0.7, 0.7)) + # origin = np.array([[-0.38063365, 0.91893092, 0.01]]) + # # end = np.array([[0, 0, 0.01]]) + # direction = np.array([[0.38173721, -0.92387953, -0.02689453]]) + # los_transform = LineOfSightTransform( + # origin[:, 0], + # origin[:, 1], + # origin[:, 2], + # direction[:, 0], + # direction[:, 1], + # direction[:, 2], + # name="", + # machine_dimensions=machine_dims, + # passes=2, + # ) los_transform.set_equilibrium(self.equilibrium) model = Interferometry(name=diag) model.set_los_transform(los_transform) @@ -403,10 +400,15 @@ def __call__(self, filepath="./results/test/", run=None, mds_write=False, pulse_ if __name__ == "__main__": run = BayesWorkflowExample( - pulse=9780, + pulse=None, phantoms=True, - diagnostics=["xrcs", "efit", "smmh1",], - opt_quantity=OPTIMISED_QUANTITY, + diagnostics=["xrcs", "efit", "smmh1","cxff_pi"], + opt_quantity=[ + "xrcs.spectra", + "cxff_pi.ti", + "efit.wp", + "smmh1.ne", + ], param_names=OPTIMISED_PARAMS, profile_params=DEFAULT_PROFILE_PARAMS, priors=DEFAULT_PRIORS, @@ -420,7 +422,7 @@ def __call__(self, filepath="./results/test/", run=None, mds_write=False, pulse_ ) run.setup_opt_data(phantoms=run.phantoms) run.setup_optimiser( - iterations=2, + iterations=100, nwalkers=50, burn_frac=0.10, sample_high_density=True, From d5ebe08c49ba9d0dad29bfd726ecee3e9aff214f Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 17 Aug 2023 11:19:32 +0100 Subject: [PATCH 138/160] moved kwargs to sample function --- indica/workflows/abstract_bayes_workflow.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/indica/workflows/abstract_bayes_workflow.py b/indica/workflows/abstract_bayes_workflow.py index 71b886ad..0f8afcfc 100644 --- a/indica/workflows/abstract_bayes_workflow.py +++ b/indica/workflows/abstract_bayes_workflow.py @@ -387,7 +387,7 @@ def _build_result_dict(self): self.result = result return self.result - def run_sampler(self): + def run_sampler(self, iterations, burn_frac): """ TODO: unsure if keeping in abstract class is best practice @@ -398,15 +398,18 @@ def run_sampler(self): result in MDSPlus node formatting """ + self.burn_frac = burn_frac + self.iterations = iterations + self.autocorr = sample_with_autocorr( self.sampler, self.start_points, - self.iterations, + iterations, self.param_names.__len__(), auto_sample=10, ) blobs = self.sampler.get_blobs( - discard=int(self.iterations * self.burn_frac), flat=True + discard=int(iterations * burn_frac), flat=True ) blob_names = self.sampler.get_blobs().flatten()[0].keys() self.samples = np.arange(0, blobs.__len__()) From 056b7f5ea8e15411b8723afac66d659a144d76f5 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 17 Aug 2023 11:53:30 +0100 Subject: [PATCH 139/160] nsamples kwarg added to sample_from_high_density_region --- indica/workflows/bayes_workflow_example.py | 25 ++++++++++++---------- 1 file changed, 14 insertions(+), 11 deletions(-) diff --git a/indica/workflows/bayes_workflow_example.py b/indica/workflows/bayes_workflow_example.py index 5d2e5868..a6a2ad0d 100644 --- a/indica/workflows/bayes_workflow_example.py +++ b/indica/workflows/bayes_workflow_example.py @@ -6,7 +6,7 @@ from indica.bayesmodels import BayesModels, get_uniform from indica.workflows.bayes_plots import plot_bayes_result from indica.workflows.abstract_bayes_workflow import AbstractBayesWorkflow -from indica.writers.bda_tree import create_nodes, write_nodes, check_analysis_run +from indica.writers.bda_tree import create_nodes, write_nodes from indica.converters.line_of_sight import LineOfSightTransform from indica.models.interferometry import Interferometry, smmh1_transform_example @@ -99,7 +99,12 @@ "Ti_prof.wcenter", "Ti_prof.peaking", ] - +OPTIMISED_QUANTITY = [ + "xrcs.spectra", + "cxff_pi.ti", + "efit.wp", + "smmh1.ne", + ] class BayesWorkflowExample(AbstractBayesWorkflow): def __init__( @@ -330,13 +335,11 @@ def _exp_data(self, **kwargs): def setup_optimiser(self, model_kwargs, - iterations=200, nwalkers=50, burn_frac=0.10, sample_high_density=False, ): self.model_kwargs = model_kwargs - self.iterations = iterations self.nwalkers = nwalkers self.burn_frac = burn_frac self.sample_high_density = sample_high_density @@ -352,7 +355,7 @@ def setup_optimiser(self, ndim = len(self.param_names) self.move = [(emcee.moves.StretchMove(), 0.9), (emcee.moves.DEMove(), 0.1)] self.sampler = emcee.EnsembleSampler( - self.nwalkers, + nwalkers, ndim, log_prob_fn=self.bayesmodel.ln_posterior, parameter_names=self.param_names, @@ -366,7 +369,7 @@ def setup_optimiser(self, def _sample_start_points(self, sample_high_density: bool = True): if sample_high_density: start_points = self.bayesmodel.sample_from_high_density_region( - self.param_names, self.sampler, self.nwalkers + self.param_names, self.sampler, self.nwalkers, nsamples=100 ) else: start_points = self.bayesmodel.sample_from_priors( @@ -374,8 +377,8 @@ def _sample_start_points(self, sample_high_density: bool = True): ) return start_points - def __call__(self, filepath="./results/test/", run=None, mds_write=False, pulse_to_write=None, plot=False, - **kwargs): + def __call__(self, filepath="./results/test/", run=None, mds_write=False, pulse_to_write=None, + plot=False, iterations=100, burn_frac=0.10, **kwargs): if mds_write: # check_analysis_run(self.pulse, self.run) self.node_structure = create_nodes( @@ -385,7 +388,7 @@ def __call__(self, filepath="./results/test/", run=None, mds_write=False, pulse_ mode="NEW", ) - self.run_sampler() + self.run_sampler(iterations=iterations, burn_frac=burn_frac) self.save_pickle(filepath=filepath) if plot: # currently requires result with DataArrays @@ -422,9 +425,7 @@ def __call__(self, filepath="./results/test/", run=None, mds_write=False, pulse_ ) run.setup_opt_data(phantoms=run.phantoms) run.setup_optimiser( - iterations=100, nwalkers=50, - burn_frac=0.10, sample_high_density=True, model_kwargs={ } @@ -435,4 +436,6 @@ def __call__(self, filepath="./results/test/", run=None, mds_write=False, pulse_ run="RUN01", mds_write=True, plot=True, + burn_frac=0.10, + iterations=100, ) From e1b35c537ac54a308b4fca77ad9f2ad4eb794e53 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 17 Aug 2023 11:53:54 +0100 Subject: [PATCH 140/160] mocked bda_tree module --- .../workflows/test_bayes_workflow_example.py | 20 +++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/tests/unit/workflows/test_bayes_workflow_example.py b/tests/unit/workflows/test_bayes_workflow_example.py index b80fc392..6dd2370c 100644 --- a/tests/unit/workflows/test_bayes_workflow_example.py +++ b/tests/unit/workflows/test_bayes_workflow_example.py @@ -1,4 +1,7 @@ import pytest +from unittest.mock import Mock +import sys + import copy import numpy as np @@ -10,11 +13,16 @@ OPTIMISED_QUANTITY, ) + + + """ TODO: Mock reader for testing experimental data reading """ +sys.modules["indica.writers.bda_tree"] = Mock() + class TestBayesWorkflowExample: def setup_class(self): @@ -39,9 +47,7 @@ def setup_class(self): model_kwargs = { "xrcs_moment_analysis": False, }, - iterations = 1, nwalkers = 20, - burn_frac = 0.10, sample_high_density = False, ) @@ -50,7 +56,13 @@ def setup_class(self): pulse_to_write=23000101, run="RUN01", mds_write = False, - plot = False,) + plot = False, + iterations=1, + burn_frac=0.10, + ) + self.sampler_settings = dict( + iterations=1, + burn_frac=0.10,) self.workflow_untouched = BayesWorkflowExample(**self.init_settings) self.workflow=None @@ -147,7 +159,7 @@ def test_worklow_has_results_after_run(self): self.workflow.setup_plasma(**self.plasma_settings) self.workflow.setup_opt_data(phantoms=True) self.workflow.setup_optimiser(**self.optimiser_settings) - self.workflow.run_sampler() + self.workflow.run_sampler(**self.sampler_settings) if not hasattr(self.workflow, "result"): raise ValueError(f"missing result in workflow object") assert True From 3733bc6f03fda0ef44198fac82323320ae5b282f Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 17 Aug 2023 11:56:36 +0100 Subject: [PATCH 141/160] fixed import --- indica/workflows/load_modelling_plasma.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/indica/workflows/load_modelling_plasma.py b/indica/workflows/load_modelling_plasma.py index e0fcd6a5..d0e0a5fe 100644 --- a/indica/workflows/load_modelling_plasma.py +++ b/indica/workflows/load_modelling_plasma.py @@ -13,7 +13,7 @@ from indica.equilibrium import Equilibrium from indica.models.charge_exchange import ChargeExchange from indica.models.diode_filters import BremsstrahlungDiode -from indica.models.helike_spectroscopy import Helike_spectroscopy +from indica.models.helike_spectroscopy import HelikeSpectrometer from indica.models.interferometry import Interferometry from indica.models.plasma import Plasma from indica.models.sxr_camera import SXRcamera From b0f3fae648bfae748ae40267ddf407c5407a08b4 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 17 Aug 2023 11:56:54 +0100 Subject: [PATCH 142/160] fixed import --- indica/workflows/load_modelling_plasma.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/indica/workflows/load_modelling_plasma.py b/indica/workflows/load_modelling_plasma.py index d0e0a5fe..957577bc 100644 --- a/indica/workflows/load_modelling_plasma.py +++ b/indica/workflows/load_modelling_plasma.py @@ -31,7 +31,7 @@ DIAGNOSTIC_MODELS = { "smmh1": Interferometry, "nirh1": Interferometry, - "xrcs": Helike_spectroscopy, + "xrcs": HelikeSpectrometer, "ts": ThomsonScattering, "cxff_pi": ChargeExchange, "cxff_tws_c": ChargeExchange, From ad9400a8053a2dbd03376214741189c7cb5a61f3 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 17 Aug 2023 11:58:12 +0100 Subject: [PATCH 143/160] black formatting --- indica/workflows/abstract_bayes_workflow.py | 128 ++++++++++--------- indica/workflows/bayes_workflow_example.py | 132 ++++++++++---------- indica/workflows/profile_workflows.py | 1 - indica/workflows/run_tomo_1d.py | 2 - indica/workflows/zeff_profile.py | 8 +- 5 files changed, 132 insertions(+), 139 deletions(-) diff --git a/indica/workflows/abstract_bayes_workflow.py b/indica/workflows/abstract_bayes_workflow.py index 0f8afcfc..a8496855 100644 --- a/indica/workflows/abstract_bayes_workflow.py +++ b/indica/workflows/abstract_bayes_workflow.py @@ -12,13 +12,13 @@ class AbstractBayesWorkflow(ABC): @abstractmethod def __init__( - self, - pulse=None, - phantoms=None, - diagnostics=None, - param_names=None, - opt_quantity=None, - priors=None, + self, + pulse=None, + phantoms=None, + diagnostics=None, + param_names=None, + opt_quantity=None, + priors=None, ): self.pulse = pulse self.phantoms = phantoms @@ -111,8 +111,8 @@ def save_phantom_profiles(self, kinetic_profiles=None): if self.phantoms: phantom_profiles = { profile_key: getattr(self.plasma, profile_key) - .sel(t=self.plasma.time_to_calculate) - .copy() + .sel(t=self.plasma.time_to_calculate) + .copy() for profile_key in kinetic_profiles } else: @@ -120,7 +120,7 @@ def save_phantom_profiles(self, kinetic_profiles=None): profile_key: getattr(self.plasma, profile_key).sel( t=self.plasma.time_to_calculate ) - * 0 + * 0 for profile_key in kinetic_profiles } @@ -225,22 +225,22 @@ def _build_result_dict(self): "RHO_TOR": self.plasma.equilibrium.rhotor.interp(t=self.plasma.t), "NE": self.blobs["electron_density"].median(dim="index"), "NI": self.blobs["ion_density"] - .sel(element=self.plasma.main_ion) - .median(dim="index"), + .sel(element=self.plasma.main_ion) + .median(dim="index"), "TE": self.blobs["electron_temperature"].median(dim="index"), "TI": self.blobs["ion_temperature"] - .sel(element=self.plasma.main_ion) - .median(dim="index"), + .sel(element=self.plasma.main_ion) + .median(dim="index"), "NFAST": self.blobs["fast_density"].median(dim="index"), "NNEUTR": self.blobs["neutral_density"].median(dim="index"), "NE_ERR": self.blobs["electron_density"].std(dim="index"), "NI_ERR": self.blobs["ion_density"] - .sel(element=self.plasma.main_ion) - .std(dim="index"), + .sel(element=self.plasma.main_ion) + .std(dim="index"), "TE_ERR": self.blobs["electron_temperature"].std(dim="index"), "TI_ERR": self.blobs["ion_temperature"] - .sel(element=self.plasma.main_ion) - .std(dim="index"), + .sel(element=self.plasma.main_ion) + .std(dim="index"), "NFAST_ERR": self.blobs["fast_density"].std(dim="index"), "NNEUTR_ERR": self.blobs["neutral_density"].std(dim="index"), } @@ -248,8 +248,8 @@ def _build_result_dict(self): **result["PROFILES"], **{ f"NIZ{num_imp + 1}": self.blobs["impurity_density"] - .sel(element=imp) - .median(dim="index") + .sel(element=imp) + .median(dim="index") for num_imp, imp in enumerate(self.plasma.impurities) }, } @@ -257,8 +257,8 @@ def _build_result_dict(self): **result["PROFILES"], **{ f"NIZ{num_imp + 1}_ERR": self.blobs["impurity_density"] - .sel(element=imp) - .std(dim="index") + .sel(element=imp) + .std(dim="index") for num_imp, imp in enumerate(self.plasma.impurities) }, } @@ -266,8 +266,8 @@ def _build_result_dict(self): **result["PROFILES"], **{ f"TIZ{num_imp + 1}": self.blobs["ion_temperature"] - .sel(element=imp) - .median(dim="index") + .sel(element=imp) + .median(dim="index") for num_imp, imp in enumerate(self.plasma.impurities) }, } @@ -275,8 +275,8 @@ def _build_result_dict(self): **result["PROFILES"], **{ f"TIZ{num_imp + 1}_ERR": self.blobs["ion_temperature"] - .sel(element=imp) - .std(dim="index") + .sel(element=imp) + .std(dim="index") for num_imp, imp in enumerate(self.plasma.impurities) }, } @@ -315,41 +315,41 @@ def _build_result_dict(self): result["GLOBAL"] = { "TI0": self.blobs["ion_temperature"] - .sel(element=self.plasma.main_ion) - .sel(rho_poloidal=0, method="nearest") - .median(dim="index"), + .sel(element=self.plasma.main_ion) + .sel(rho_poloidal=0, method="nearest") + .median(dim="index"), "TE0": self.blobs["electron_temperature"] - .sel(rho_poloidal=0, method="nearest") - .median(dim="index"), + .sel(rho_poloidal=0, method="nearest") + .median(dim="index"), "NE0": self.blobs["electron_density"] - .sel(rho_poloidal=0, method="nearest") - .median(dim="index"), + .sel(rho_poloidal=0, method="nearest") + .median(dim="index"), "NI0": self.blobs["ion_density"] - .sel(element=self.plasma.main_ion) - .sel(rho_poloidal=0, method="nearest") - .median(dim="index"), + .sel(element=self.plasma.main_ion) + .sel(rho_poloidal=0, method="nearest") + .median(dim="index"), "TI0_ERR": self.blobs["ion_temperature"] - .sel(element=self.plasma.main_ion) - .sel(rho_poloidal=0, method="nearest") - .std(dim="index"), + .sel(element=self.plasma.main_ion) + .sel(rho_poloidal=0, method="nearest") + .std(dim="index"), "TE0_ERR": self.blobs["electron_temperature"] - .sel(rho_poloidal=0, method="nearest") - .std(dim="index"), + .sel(rho_poloidal=0, method="nearest") + .std(dim="index"), "NE0_ERR": self.blobs["electron_density"] - .sel(rho_poloidal=0, method="nearest") - .std(dim="index"), + .sel(rho_poloidal=0, method="nearest") + .std(dim="index"), "NI0_ERR": self.blobs["ion_density"] - .sel(element=self.plasma.main_ion) - .sel(rho_poloidal=0, method="nearest") - .std(dim="index"), + .sel(element=self.plasma.main_ion) + .sel(rho_poloidal=0, method="nearest") + .std(dim="index"), } result["GLOBAL"] = { **result["GLOBAL"], **{ f"TI0Z{num_imp + 1}": self.blobs["ion_temperature"] - .sel(element=imp) - .sel(rho_poloidal=0, method="nearest") - .median(dim="index") + .sel(element=imp) + .sel(rho_poloidal=0, method="nearest") + .median(dim="index") for num_imp, imp in enumerate(self.plasma.impurities) }, } @@ -357,9 +357,9 @@ def _build_result_dict(self): **result["GLOBAL"], **{ f"TI0Z{num_imp + 1}_ERR": self.blobs["ion_temperature"] - .sel(element=imp) - .sel(rho_poloidal=0, method="nearest") - .std(dim="index") + .sel(element=imp) + .sel(rho_poloidal=0, method="nearest") + .std(dim="index") for num_imp, imp in enumerate(self.plasma.impurities) }, } @@ -367,9 +367,9 @@ def _build_result_dict(self): **result["GLOBAL"], **{ f"NI0Z{num_imp + 1}": self.blobs["impurity_density"] - .sel(element=imp) - .sel(rho_poloidal=0, method="nearest") - .median(dim="index") + .sel(element=imp) + .sel(rho_poloidal=0, method="nearest") + .median(dim="index") for num_imp, imp in enumerate(self.plasma.impurities) }, } @@ -377,9 +377,9 @@ def _build_result_dict(self): **result["GLOBAL"], **{ f"NI0Z{num_imp + 1}_ERR": self.blobs["impurity_density"] - .sel(element=imp) - .sel(rho_poloidal=0, method="nearest") - .std(dim="index") + .sel(element=imp) + .sel(rho_poloidal=0, method="nearest") + .std(dim="index") for num_imp, imp in enumerate(self.plasma.impurities) }, } @@ -408,9 +408,7 @@ def run_sampler(self, iterations, burn_frac): self.param_names.__len__(), auto_sample=10, ) - blobs = self.sampler.get_blobs( - discard=int(iterations * burn_frac), flat=True - ) + blobs = self.sampler.get_blobs(discard=int(iterations * burn_frac), flat=True) blob_names = self.sampler.get_blobs().flatten()[0].keys() self.samples = np.arange(0, blobs.__len__()) @@ -459,9 +457,9 @@ def sample_with_autocorr(sampler, start_points, iterations, n_params, auto_sampl autocorr = np.ones(shape=(iterations, n_params)) * np.nan old_tau = np.inf for sample in sampler.sample( - start_points, - iterations=iterations, - progress=True, + start_points, + iterations=iterations, + progress=True, ): if sampler.iteration % auto_sample: continue @@ -472,5 +470,5 @@ def sample_with_autocorr(sampler, start_points, iterations, n_params, auto_sampl if converged: break old_tau = new_tau - autocorr = autocorr[: sampler.iteration, ] + autocorr = autocorr[: sampler.iteration,] return autocorr diff --git a/indica/workflows/bayes_workflow_example.py b/indica/workflows/bayes_workflow_example.py index a6a2ad0d..eebc3d82 100644 --- a/indica/workflows/bayes_workflow_example.py +++ b/indica/workflows/bayes_workflow_example.py @@ -50,7 +50,6 @@ "Ne_prof.wped": get_uniform(1, 6), "Ne_prof.wcenter": get_uniform(0.1, 0.8), "Ne_prof.peaking": get_uniform(1, 6), - "Nimp_prof.y0": loguniform(1e16, 1e18), "Nimp_prof.y1": get_uniform(1e15, 1e16), "Ne_prof.y0/Nimp_prof.y0": lambda x1, x2: np.where( @@ -63,12 +62,10 @@ "Nimp_prof.peaking/Ne_prof.peaking": lambda x1, x2: np.where( (x1 > x2), 1, 0 ), # impurity always more peaked - "Te_prof.y0": get_uniform(1000, 5000), "Te_prof.wped": get_uniform(1, 6), "Te_prof.wcenter": get_uniform(0.1, 0.8), "Te_prof.peaking": get_uniform(1, 10), - # "Ti_prof.y0/Te_prof.y0": lambda x1, x2: np.where(x1 > x2, 1, 0), # hot ion mode "Ti_prof.y0": get_uniform(1000, 10000), "Ti_prof.wped": get_uniform(1, 6), @@ -82,43 +79,41 @@ # "Ne_prof.peaking", # "Ne_prof.wcenter", # "Ne_prof.wped", - # "Nimp_prof.y1", "Nimp_prof.y0", # "Nimp_prof.wcenter", # "Nimp_prof.wped", # "Nimp_prof.peaking", - "Te_prof.y0", # "Te_prof.wped", "Te_prof.wcenter", "Te_prof.peaking", - "Ti_prof.y0", # "Ti_prof.wped", "Ti_prof.wcenter", "Ti_prof.peaking", ] OPTIMISED_QUANTITY = [ - "xrcs.spectra", - "cxff_pi.ti", - "efit.wp", - "smmh1.ne", - ] + "xrcs.spectra", + "cxff_pi.ti", + "efit.wp", + "smmh1.ne", +] + class BayesWorkflowExample(AbstractBayesWorkflow): def __init__( - self, - pulse: int = None, - diagnostics: list = None, - param_names: list = None, - opt_quantity: list = None, - priors: dict = None, - profile_params: dict = None, - phantoms: bool = False, - tstart=0.02, - tend=0.10, - dt=0.005, + self, + pulse: int = None, + diagnostics: list = None, + param_names: list = None, + opt_quantity: list = None, + priors: dict = None, + profile_params: dict = None, + phantoms: bool = False, + tstart=0.02, + tend=0.10, + dt=0.005, ): self.pulse = pulse self.diagnostics = diagnostics @@ -160,18 +155,18 @@ def __init__( self.read_data(self.diagnostics) self.setup_models(self.diagnostics) - def setup_plasma(self, - tstart=None, - tend=None, - dt=None, - tsample=0.050, - main_ion = "h", - impurities = ("ar", "c"), - impurity_concentration = (0.001, 0.04), - n_rad=20, - **kwargs - ): - + def setup_plasma( + self, + tstart=None, + tend=None, + dt=None, + tsample=0.050, + main_ion="h", + impurities=("ar", "c"), + impurity_concentration=(0.001, 0.04), + n_rad=20, + **kwargs, + ): if not all([tstart, tend, dt]): tstart = self.tstart tend = self.tend @@ -227,9 +222,7 @@ def setup_models(self, diagnostics: list): window = None if hasattr(self, "data"): if diag in self.data.keys(): - window = ( - self.data[diag]["spectra"].wavelength.values - ) + window = self.data[diag]["spectra"].wavelength.values model = HelikeSpectrometer( name="xrcs", @@ -266,28 +259,28 @@ def _phantom_data(self, noise=False, noise_factor=0.1, **kwargs): if "smmh1" in self.diagnostics: opt_data["smmh1.ne"] = ( self.models["smmh1"]() - .pop("ne") - .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) + .pop("ne") + .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) ) if "xrcs" in self.diagnostics: opt_data["xrcs.spectra"] = ( self.models["xrcs"]() - .pop("spectra") - .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) + .pop("spectra") + .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) ) opt_data["xrcs.spectra"]["error"] = np.sqrt(opt_data["xrcs.spectra"]) if "cxff_pi" in self.diagnostics: cxrs_data = ( self.models["cxff_pi"]() - .pop("ti") - .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) + .pop("ti") + .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) ) opt_data["cxff_pi.ti"] = cxrs_data.where(cxrs_data.channel == 2, drop=True) if "efit" in self.diagnostics: opt_data["efit.wp"] = ( self.models["efit"]() - .pop("wp") - .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) + .pop("wp") + .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) ) if noise: @@ -313,7 +306,7 @@ def _exp_data(self, **kwargs): opt_data = flatdict.FlatDict(self.data, ".") if "xrcs" in self.diagnostics: opt_data["xrcs.spectra"]["wavelength"] = ( - opt_data["xrcs.spectra"].wavelength * 0.1 + opt_data["xrcs.spectra"].wavelength * 0.1 ) background = opt_data["xrcs.spectra"].where( (opt_data["xrcs.spectra"].wavelength < 0.392) @@ -333,12 +326,13 @@ def _exp_data(self, **kwargs): ) return opt_data - def setup_optimiser(self, - model_kwargs, - nwalkers=50, - burn_frac=0.10, - sample_high_density=False, - ): + def setup_optimiser( + self, + model_kwargs, + nwalkers=50, + burn_frac=0.10, + sample_high_density=False, + ): self.model_kwargs = model_kwargs self.nwalkers = nwalkers self.burn_frac = burn_frac @@ -377,8 +371,17 @@ def _sample_start_points(self, sample_high_density: bool = True): ) return start_points - def __call__(self, filepath="./results/test/", run=None, mds_write=False, pulse_to_write=None, - plot=False, iterations=100, burn_frac=0.10, **kwargs): + def __call__( + self, + filepath="./results/test/", + run=None, + mds_write=False, + pulse_to_write=None, + plot=False, + iterations=100, + burn_frac=0.10, + **kwargs, + ): if mds_write: # check_analysis_run(self.pulse, self.run) self.node_structure = create_nodes( @@ -405,13 +408,13 @@ def __call__(self, filepath="./results/test/", run=None, mds_write=False, pulse_ run = BayesWorkflowExample( pulse=None, phantoms=True, - diagnostics=["xrcs", "efit", "smmh1","cxff_pi"], + diagnostics=["xrcs", "efit", "smmh1", "cxff_pi"], opt_quantity=[ - "xrcs.spectra", - "cxff_pi.ti", - "efit.wp", - "smmh1.ne", - ], + "xrcs.spectra", + "cxff_pi.ti", + "efit.wp", + "smmh1.ne", + ], param_names=OPTIMISED_PARAMS, profile_params=DEFAULT_PROFILE_PARAMS, priors=DEFAULT_PRIORS, @@ -421,15 +424,10 @@ def __call__(self, filepath="./results/test/", run=None, mds_write=False, pulse_ ) run.setup_plasma( - tsample=0.060, + tsample=0.060, ) run.setup_opt_data(phantoms=run.phantoms) - run.setup_optimiser( - nwalkers=50, - sample_high_density=True, - model_kwargs={ - } - ) + run.setup_optimiser(nwalkers=50, sample_high_density=True, model_kwargs={}) results = run( filepath="./results/test/", pulse_to_write=23000101, diff --git a/indica/workflows/profile_workflows.py b/indica/workflows/profile_workflows.py index bf072059..27c4e292 100644 --- a/indica/workflows/profile_workflows.py +++ b/indica/workflows/profile_workflows.py @@ -60,7 +60,6 @@ def profile_scans_pca( save_fig: bool = False, fig_path: str = None, ): - if fig_path is None: fig_path = FIG_PATH diff --git a/indica/workflows/run_tomo_1d.py b/indica/workflows/run_tomo_1d.py index 5398c5db..c4a73fa3 100644 --- a/indica/workflows/run_tomo_1d.py +++ b/indica/workflows/run_tomo_1d.py @@ -201,7 +201,6 @@ def sxrc_xy( save_fig: bool = True, instrument="sxrc_xy2", ): - if input_dict is None: st40 = ReadST40(pulse, tstart, tend, dt=dt) st40(instruments=[instrument], map_diagnostics=False) @@ -298,7 +297,6 @@ def old_camera( reg_level_guess: float = 0.5, input_dict: dict = None, ): - if input_dict is None: st40 = ReadST40(pulse, tstart, tend) # return st40 diff --git a/indica/workflows/zeff_profile.py b/indica/workflows/zeff_profile.py index dd88f85d..a87ac99b 100644 --- a/indica/workflows/zeff_profile.py +++ b/indica/workflows/zeff_profile.py @@ -196,7 +196,10 @@ def prepare_data_cxrs( continue data = st40.binned_data[instrument] attrs = data["spectra"].attrs - (data["background"], data["brightness"],) = models[ + ( + data["background"], + data["brightness"], + ) = models[ instrument ].integrate_spectra(data["spectra"]) data["background"].attrs = attrs @@ -222,7 +225,6 @@ def prepare_inputs( phantom_data: bool = True, ts_side: str = "LFS", ): - flat_data: dict = {} models: dict = {} @@ -350,7 +352,6 @@ def run_bayes( phantom_data: bool = True, ts_side: str = "LFS", ): - plasma, models, flat_data, input_profiles = prepare_inputs( pulse, tstart=tstart, @@ -430,7 +431,6 @@ def run_inversion( phantom_data: bool = True, ts_side: str = "LFS", ): - plasma, models, flat_data, input_profiles = prepare_inputs( pulse, tstart=tstart, From f0a1d658ff9fbc635f5750048a9cb79343e331f1 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 17 Aug 2023 11:59:18 +0100 Subject: [PATCH 144/160] black formatting --- .../workflows/test_bayes_workflow_example.py | 49 ++++++++++--------- 1 file changed, 25 insertions(+), 24 deletions(-) diff --git a/tests/unit/workflows/test_bayes_workflow_example.py b/tests/unit/workflows/test_bayes_workflow_example.py index 6dd2370c..a389d381 100644 --- a/tests/unit/workflows/test_bayes_workflow_example.py +++ b/tests/unit/workflows/test_bayes_workflow_example.py @@ -14,8 +14,6 @@ ) - - """ TODO: Mock reader for testing experimental data reading @@ -28,7 +26,7 @@ class TestBayesWorkflowExample: def setup_class(self): self.init_settings = dict( pulse=None, - phantoms = True, + phantoms=True, diagnostics=["xrcs", "efit", "smmh1", "cxff_pi"], opt_quantity=OPTIMISED_QUANTITY, param_names=OPTIMISED_PARAMS, @@ -39,33 +37,33 @@ def setup_class(self): dt=0.005, ) self.plasma_settings = dict( - tsample=0.060, + tsample=0.060, ) - self.optimiser_settings = dict( - model_kwargs = { - "xrcs_moment_analysis": False, - }, - nwalkers = 20, - sample_high_density = False, + model_kwargs={ + "xrcs_moment_analysis": False, + }, + nwalkers=20, + sample_high_density=False, ) self.call_settings = dict( filepath=None, pulse_to_write=23000101, run="RUN01", - mds_write = False, - plot = False, + mds_write=False, + plot=False, iterations=1, burn_frac=0.10, ) self.sampler_settings = dict( iterations=1, - burn_frac=0.10,) + burn_frac=0.10, + ) self.workflow_untouched = BayesWorkflowExample(**self.init_settings) - self.workflow=None + self.workflow = None def setup_method(self): self.workflow = copy.deepcopy(self.workflow_untouched) @@ -74,12 +72,7 @@ def teardown_method(self): self.workflow = None def test_workflow_initializes(self): - attributes_to_check = [ - "data", - "reader", - "models", - "equilibrium" - ] + attributes_to_check = ["data", "reader", "models", "equilibrium"] for attribute in attributes_to_check: if not hasattr(self.workflow, attribute): raise ValueError(f"missing {attribute} in workflow object") @@ -145,14 +138,22 @@ def test_phantom_data_runs_with_noise_added(self): def test_sampling_from_priors(self): self.workflow.setup_plasma(**self.plasma_settings) - self.workflow.setup_opt_data(phantoms=True, ) - self.workflow.setup_optimiser(**dict(self.optimiser_settings, **{"sample_high_density":False})) + self.workflow.setup_opt_data( + phantoms=True, + ) + self.workflow.setup_optimiser( + **dict(self.optimiser_settings, **{"sample_high_density": False}) + ) assert True def test_sampling_from_high_density(self): self.workflow.setup_plasma(**self.plasma_settings) - self.workflow.setup_opt_data(phantoms=True, ) - self.workflow.setup_optimiser(**dict(self.optimiser_settings, **{"sample_high_density":True})) + self.workflow.setup_opt_data( + phantoms=True, + ) + self.workflow.setup_optimiser( + **dict(self.optimiser_settings, **{"sample_high_density": True}) + ) assert True def test_worklow_has_results_after_run(self): From 3b1dcd8d4d9865a6aedf18777f7606889c2e01bc Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 17 Aug 2023 13:18:57 +0100 Subject: [PATCH 145/160] fixed where mocked module is imported --- tests/unit/workflows/test_bayes_workflow_example.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/unit/workflows/test_bayes_workflow_example.py b/tests/unit/workflows/test_bayes_workflow_example.py index a389d381..6db6e361 100644 --- a/tests/unit/workflows/test_bayes_workflow_example.py +++ b/tests/unit/workflows/test_bayes_workflow_example.py @@ -4,6 +4,7 @@ import copy import numpy as np +sys.modules["indica.writers.bda_tree"] = Mock() from indica.workflows.bayes_workflow_example import ( BayesWorkflowExample, @@ -19,8 +20,6 @@ Mock reader for testing experimental data reading """ -sys.modules["indica.writers.bda_tree"] = Mock() - class TestBayesWorkflowExample: def setup_class(self): From b8b6ceebca889385bc85366d7d79ce9fe33e9ada Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 17 Aug 2023 13:27:36 +0100 Subject: [PATCH 146/160] black formatting --- indica/models/helike_spectroscopy.py | 2 ++ indica/models/interferometry.py | 5 ++--- indica/models/plasma.py | 8 +++++++- tests/unit/workflows/test_bayes_workflow_example.py | 1 + 4 files changed, 12 insertions(+), 4 deletions(-) diff --git a/indica/models/helike_spectroscopy.py b/indica/models/helike_spectroscopy.py index 5fd32aae..90ac5c4e 100644 --- a/indica/models/helike_spectroscopy.py +++ b/indica/models/helike_spectroscopy.py @@ -465,6 +465,7 @@ def __call__( self._build_bckc_dictionary() return self.bckc + def helike_transform_example(nchannels): los_end = np.full((nchannels, 3), 0.0) los_end[:, 0] = 0.17 @@ -488,6 +489,7 @@ def helike_transform_example(nchannels): ) return los_transform + def example_run( pulse: int = None, plasma=None, plot=False, moment_analysis: bool = False, **kwargs ): diff --git a/indica/models/interferometry.py b/indica/models/interferometry.py index c9354205..5a7cf49b 100644 --- a/indica/models/interferometry.py +++ b/indica/models/interferometry.py @@ -90,9 +90,8 @@ def __call__( def smmh1_transform_example(nchannels): - - los_start = np.array([[0.8, 0, 0]]) * np.ones((nchannels,3)) - los_start[:,2] = np.linspace(0, -0.2, nchannels) + los_start = np.array([[0.8, 0, 0]]) * np.ones((nchannels, 3)) + los_start[:, 2] = np.linspace(0, -0.2, nchannels) los_end = np.array([[0.17, 0, 0]]) * np.ones((nchannels, 3)) los_end[:, 2] = np.linspace(0, -0.2, nchannels) origin = los_start diff --git a/indica/models/plasma.py b/indica/models/plasma.py index 86630796..435b2d1d 100644 --- a/indica/models/plasma.py +++ b/indica/models/plasma.py @@ -517,7 +517,13 @@ def update_profiles( else: raise ValueError(f"parameter: {key} not found in {prefix}") - for key in ["electron_density", "electron_temperature", "ion_temperature", "toroidal_rotation", "impurity_density"]: + for key in [ + "electron_density", + "electron_temperature", + "ion_temperature", + "toroidal_rotation", + "impurity_density", + ]: self.assign_profiles(key, t=self.time_to_calculate) @property diff --git a/tests/unit/workflows/test_bayes_workflow_example.py b/tests/unit/workflows/test_bayes_workflow_example.py index 6db6e361..faaa9347 100644 --- a/tests/unit/workflows/test_bayes_workflow_example.py +++ b/tests/unit/workflows/test_bayes_workflow_example.py @@ -4,6 +4,7 @@ import copy import numpy as np + sys.modules["indica.writers.bda_tree"] = Mock() from indica.workflows.bayes_workflow_example import ( From d3e4e5c4779cf80a20c7e657c35638cdab118a91 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 17 Aug 2023 14:00:23 +0100 Subject: [PATCH 147/160] mocking bda_tree import --- indica/bayesmodels.py | 4 +++- tests/conftest.py | 1 + 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/indica/bayesmodels.py b/indica/bayesmodels.py index 47685db7..900946a6 100644 --- a/indica/bayesmodels.py +++ b/indica/bayesmodels.py @@ -1,5 +1,6 @@ from copy import deepcopy import warnings + import numpy as np from scipy.stats import uniform @@ -82,7 +83,8 @@ def _build_bckc(self, params, **kwargs): """ self.bckc: dict = {} for model in self.diagnostic_models: - # removes "model.name." from params and kwargs then passes them to model e.g. xrcs.background -> background + # removes "model.name." from params and kwargs then passes them to model + # e.g. xrcs.background -> background _nuisance_params = { param_name.replace(model.name + ".", ""): param_value for param_name, param_value in params.items() diff --git a/tests/conftest.py b/tests/conftest.py index 209c8a1b..e5b9475a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -25,3 +25,4 @@ # Turn off import of modules that cnnot be installed in CI sys.modules["indica.readers.st40reader"] = mock.MagicMock() sys.modules["indica.readers.st40reader.ST40Reader"] = mock.MagicMock() +sys.modules["indica.writers.bda_tree"] = mock.Mock() From a7ec4cdf724588810c4a2f6984c5f5e6c80455bf Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 17 Aug 2023 15:26:39 +0100 Subject: [PATCH 148/160] precommit --- indica/models/plasma.py | 1 - indica/workflows/abstract_bayes_workflow.py | 22 +++++++----- indica/workflows/bayes_plots.py | 11 +++--- indica/workflows/bayes_workflow_example.py | 32 +++++++++-------- indica/workflows/zeff_profile.py | 5 +-- indica/writers/bda_tree.py | 9 +++-- .../workflows/test_bayes_workflow_example.py | 35 +++++++------------ 7 files changed, 56 insertions(+), 59 deletions(-) diff --git a/indica/models/plasma.py b/indica/models/plasma.py index 435b2d1d..0c635f88 100644 --- a/indica/models/plasma.py +++ b/indica/models/plasma.py @@ -14,7 +14,6 @@ from indica.datatypes import ELEMENTS from indica.equilibrium import Equilibrium from indica.equilibrium import fake_equilibrium_data -from indica.numpy_typing import LabeledArray from indica.operators.atomic_data import FractionalAbundance from indica.operators.atomic_data import PowerLoss import indica.physics as ph diff --git a/indica/workflows/abstract_bayes_workflow.py b/indica/workflows/abstract_bayes_workflow.py index a8496855..f62c037d 100644 --- a/indica/workflows/abstract_bayes_workflow.py +++ b/indica/workflows/abstract_bayes_workflow.py @@ -1,12 +1,14 @@ -from abc import ABC, abstractmethod -import numpy as np -import xarray as xr -import pandas as pd +from abc import ABC +from abc import abstractmethod from pathlib import Path import pickle -from indica.readers.read_st40 import ReadST40 +import numpy as np +import pandas as pd +import xarray as xr + from indica.equilibrium import fake_equilibrium +from indica.readers.read_st40 import ReadST40 class AbstractBayesWorkflow(ABC): @@ -162,7 +164,7 @@ def _build_result_dict(self): } result["INPUT"]["WORKFLOW"] = { diag_name.upper(): { - "PULSE": self.pulse, # Change this if different pulses used for diagnostics + "PULSE": self.pulse, # Change this if different pulses used "USAGE": "".join( [quantity[1] for quantity in quant_list if quantity[0] == diag_name] ), @@ -464,11 +466,15 @@ def sample_with_autocorr(sampler, start_points, iterations, n_params, auto_sampl if sampler.iteration % auto_sample: continue new_tau = sampler.get_autocorr_time(tol=0) - autocorr[sampler.iteration - 1,] = new_tau + autocorr[ + sampler.iteration - 1, + ] = new_tau converged = np.all(new_tau * 50 < sampler.iteration) converged &= np.all(np.abs(old_tau - new_tau) / new_tau < 0.01) if converged: break old_tau = new_tau - autocorr = autocorr[: sampler.iteration,] + autocorr = autocorr[ + : sampler.iteration, + ] return autocorr diff --git a/indica/workflows/bayes_plots.py b/indica/workflows/bayes_plots.py index c027c004..93955202 100644 --- a/indica/workflows/bayes_plots.py +++ b/indica/workflows/bayes_plots.py @@ -1,12 +1,13 @@ import os +import pickle import corner +import flatdict import matplotlib.pyplot as plt import numpy as np -import pickle -import flatdict -from indica.utilities import set_plot_rcparams, set_axis_sci +from indica.utilities import set_axis_sci +from indica.utilities import set_plot_rcparams def plot_profile( @@ -406,7 +407,7 @@ def plot_bayes_result( profiles[key], key, figheader=figheader, - filename=f"ar density", + filename="ar density", filetype=filetype, phantom_profile=phantom_profiles, color="red", @@ -417,7 +418,7 @@ def plot_bayes_result( profiles[key], key, figheader=figheader, - filename=f"c density", + filename="c density", filetype=filetype, phantom_profile=phantom_profiles, color="red", diff --git a/indica/workflows/bayes_workflow_example.py b/indica/workflows/bayes_workflow_example.py index eebc3d82..53809312 100644 --- a/indica/workflows/bayes_workflow_example.py +++ b/indica/workflows/bayes_workflow_example.py @@ -1,22 +1,22 @@ import emcee -import numpy as np import flatdict +import numpy as np from scipy.stats import loguniform -from indica.bayesmodels import BayesModels, get_uniform -from indica.workflows.bayes_plots import plot_bayes_result -from indica.workflows.abstract_bayes_workflow import AbstractBayesWorkflow -from indica.writers.bda_tree import create_nodes, write_nodes -from indica.converters.line_of_sight import LineOfSightTransform - -from indica.models.interferometry import Interferometry, smmh1_transform_example -from indica.models.helike_spectroscopy import ( - HelikeSpectrometer, - helike_transform_example, -) -from indica.models.charge_exchange import ChargeExchange, pi_transform_example +from indica.bayesmodels import BayesModels +from indica.bayesmodels import get_uniform +from indica.models.charge_exchange import ChargeExchange +from indica.models.charge_exchange import pi_transform_example from indica.models.equilibrium_reconstruction import EquilibriumReconstruction +from indica.models.helike_spectroscopy import helike_transform_example +from indica.models.helike_spectroscopy import HelikeSpectrometer +from indica.models.interferometry import Interferometry +from indica.models.interferometry import smmh1_transform_example from indica.models.plasma import Plasma +from indica.workflows.abstract_bayes_workflow import AbstractBayesWorkflow +from indica.workflows.bayes_plots import plot_bayes_result +from indica.writers.bda_tree import create_nodes +from indica.writers.bda_tree import write_nodes # global configurations DEFAULT_PROFILE_PARAMS = { @@ -289,7 +289,11 @@ def _phantom_data(self, noise=False, noise_factor=0.1, **kwargs): ].max().values * np.random.normal(0, noise_factor, None) opt_data["xrcs.spectra"] = opt_data["xrcs.spectra"] + np.random.normal( 0, - np.sqrt(opt_data["xrcs.spectra"].values[0,]), + np.sqrt( + opt_data["xrcs.spectra"].values[ + 0, + ] + ), opt_data["xrcs.spectra"].shape[1], ) opt_data["cxff_pi.ti"] = opt_data["cxff_pi.ti"] + opt_data[ diff --git a/indica/workflows/zeff_profile.py b/indica/workflows/zeff_profile.py index a87ac99b..adf29bda 100644 --- a/indica/workflows/zeff_profile.py +++ b/indica/workflows/zeff_profile.py @@ -196,10 +196,7 @@ def prepare_data_cxrs( continue data = st40.binned_data[instrument] attrs = data["spectra"].attrs - ( - data["background"], - data["brightness"], - ) = models[ + (data["background"], data["brightness"],) = models[ instrument ].integrate_spectra(data["spectra"]) data["background"].attrs = attrs diff --git a/indica/writers/bda_tree.py b/indica/writers/bda_tree.py index 6f25ecb7..5a5878a8 100644 --- a/indica/writers/bda_tree.py +++ b/indica/writers/bda_tree.py @@ -1,9 +1,10 @@ # Create trees from ppac standard utility tools -import standard_utility as util -from MDSplus import Connection import sys + +from MDSplus import Connection import numpy as np +import standard_utility as util def bda(): @@ -165,8 +166,6 @@ def create_nodes( } for diag_name in diag_names } - # for diag_name in diag_names: - # diag_nodes[diag_name]["RUN"] = ("TEXT", f"RUN from which {diag_name} data was taken") nodes = { "RUN": ("TEXT", "RUN used for diagnostic"), @@ -219,7 +218,7 @@ def create_nodes( def write_nodes(pulse, node_info, data): util.StandardNodeWriting( - pulse_number=pulse, # pulse number for which data should be written to the structure + pulse_number=pulse, # pulse number for which data should be written dict_node_info=node_info, # node information file nodes_to_write=[], # selective nodes to be written data_to_write=data, diff --git a/tests/unit/workflows/test_bayes_workflow_example.py b/tests/unit/workflows/test_bayes_workflow_example.py index faaa9347..d3111cb1 100644 --- a/tests/unit/workflows/test_bayes_workflow_example.py +++ b/tests/unit/workflows/test_bayes_workflow_example.py @@ -1,20 +1,14 @@ -import pytest -from unittest.mock import Mock -import sys - import copy -import numpy as np - -sys.modules["indica.writers.bda_tree"] = Mock() +import sys +from unittest.mock import Mock -from indica.workflows.bayes_workflow_example import ( - BayesWorkflowExample, - DEFAULT_PRIORS, - DEFAULT_PROFILE_PARAMS, - OPTIMISED_PARAMS, - OPTIMISED_QUANTITY, -) +import pytest +from indica.workflows.bayes_workflow_example import BayesWorkflowExample +from indica.workflows.bayes_workflow_example import DEFAULT_PRIORS +from indica.workflows.bayes_workflow_example import DEFAULT_PROFILE_PARAMS +from indica.workflows.bayes_workflow_example import OPTIMISED_PARAMS +from indica.workflows.bayes_workflow_example import OPTIMISED_QUANTITY """ TODO: @@ -80,18 +74,15 @@ def test_workflow_initializes(self): def test_init_phantoms_false_with_example_plasma(self): with pytest.raises(ValueError): - example = BayesWorkflowExample( - dict(self.init_settings, **{"phantoms": False}) - ) + BayesWorkflowExample(dict(self.init_settings, **{"phantoms": False})) def test_init_not_including_all_required_inputs(self): with pytest.raises(ValueError): - example = BayesWorkflowExample( - dict(self.init_settings, **{"param_names": None}) - ) + BayesWorkflowExample(dict(self.init_settings, **{"param_names": None})) # def test_reader_has_read_all_diagnostic_data(self): - # assert all(diag_name in self.workflow.reader.keys() for diag_name in self.workflow.diagnostics) + # assert all(diag_name in self.workflow.reader.keys() + # for diag_name in self.workflow.diagnostics) def test_plasma_has_equilibrium(self): self.workflow.setup_plasma(**self.plasma_settings) @@ -162,7 +153,7 @@ def test_worklow_has_results_after_run(self): self.workflow.setup_optimiser(**self.optimiser_settings) self.workflow.run_sampler(**self.sampler_settings) if not hasattr(self.workflow, "result"): - raise ValueError(f"missing result in workflow object") + raise ValueError("missing result in workflow object") assert True From b1babc4ad83f6b23a60f2cffb75f2d49b889fe16 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 17 Aug 2023 15:27:43 +0100 Subject: [PATCH 149/160] precommit --- tests/unit/workflows/test_bayes_workflow_example.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/unit/workflows/test_bayes_workflow_example.py b/tests/unit/workflows/test_bayes_workflow_example.py index d3111cb1..d725eb47 100644 --- a/tests/unit/workflows/test_bayes_workflow_example.py +++ b/tests/unit/workflows/test_bayes_workflow_example.py @@ -1,6 +1,4 @@ import copy -import sys -from unittest.mock import Mock import pytest From 8e2f2d8082a5f0d6e4763ceb22d20f07936769cd Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 17 Aug 2023 15:38:35 +0100 Subject: [PATCH 150/160] moving read_data tvector to kwargs --- indica/workflows/abstract_bayes_workflow.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/indica/workflows/abstract_bayes_workflow.py b/indica/workflows/abstract_bayes_workflow.py index f62c037d..4b7d2ef8 100644 --- a/indica/workflows/abstract_bayes_workflow.py +++ b/indica/workflows/abstract_bayes_workflow.py @@ -31,19 +31,19 @@ def __init__( self.read_data(self.diagnostics) self.setup_models(self.diagnostics) - def read_test_data(self, diagnostic_transforms: dict): + def read_test_data(self, diagnostic_transforms: dict, tstart=None, tend=None, dt=None): # Used with phantom data for purposes of tests print("Reading fake data") self.reader = None self.equilibrium = fake_equilibrium( - self.tstart, - self.tend, - self.dt, + tstart, + tend, + dt, ) self.transforms = diagnostic_transforms self.data = {} - def read_data(self, diagnostics: list): + def read_data(self, diagnostics: list, tstart=None, tend=None, dt=None): self.reader = ReadST40( self.pulse, tstart=self.tstart, tend=self.tend, dt=self.dt ) From 519c4eda6a8b784586ebe72d26e3f36c2a36ec80 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 17 Aug 2023 15:45:01 +0100 Subject: [PATCH 151/160] precommit formatting --- indica/bayesmodels.py | 2 +- indica/workflows/abstract_bayes_workflow.py | 5 ++--- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/indica/bayesmodels.py b/indica/bayesmodels.py index 900946a6..a05e8c77 100644 --- a/indica/bayesmodels.py +++ b/indica/bayesmodels.py @@ -194,7 +194,7 @@ def sample_from_priors(self, param_names, size=10): return samples.transpose() def sample_from_high_density_region( - self, param_names: list, sampler: object, nwalkers: int, nsamples=100 + self, param_names: list, sampler, nwalkers: int, nsamples=100 ): start_points = self.sample_from_priors(param_names, size=nsamples) diff --git a/indica/workflows/abstract_bayes_workflow.py b/indica/workflows/abstract_bayes_workflow.py index 4b7d2ef8..3e83fd4e 100644 --- a/indica/workflows/abstract_bayes_workflow.py +++ b/indica/workflows/abstract_bayes_workflow.py @@ -34,14 +34,13 @@ def __init__( def read_test_data(self, diagnostic_transforms: dict, tstart=None, tend=None, dt=None): # Used with phantom data for purposes of tests print("Reading fake data") - self.reader = None self.equilibrium = fake_equilibrium( tstart, tend, dt, ) self.transforms = diagnostic_transforms - self.data = {} + self.data: dict = {} def read_data(self, diagnostics: list, tstart=None, tend=None, dt=None): self.reader = ReadST40( @@ -67,7 +66,7 @@ def setup_models(self, diagnostics: list): Initialising models normally requires data to be read so transforms can be set """ - self.models = {} + self.models: dict = {} @abstractmethod def _phantom_data(self): From f7623ceebceace753bbf5738d66af10c100a5d01 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 17 Aug 2023 15:50:53 +0100 Subject: [PATCH 152/160] precommit formatting --- indica/workflows/abstract_bayes_workflow.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/indica/workflows/abstract_bayes_workflow.py b/indica/workflows/abstract_bayes_workflow.py index 3e83fd4e..3221e186 100644 --- a/indica/workflows/abstract_bayes_workflow.py +++ b/indica/workflows/abstract_bayes_workflow.py @@ -31,7 +31,9 @@ def __init__( self.read_data(self.diagnostics) self.setup_models(self.diagnostics) - def read_test_data(self, diagnostic_transforms: dict, tstart=None, tend=None, dt=None): + def read_test_data( + self, diagnostic_transforms: dict, tstart=None, tend=None, dt=None + ): # Used with phantom data for purposes of tests print("Reading fake data") self.equilibrium = fake_equilibrium( @@ -44,7 +46,7 @@ def read_test_data(self, diagnostic_transforms: dict, tstart=None, tend=None, dt def read_data(self, diagnostics: list, tstart=None, tend=None, dt=None): self.reader = ReadST40( - self.pulse, tstart=self.tstart, tend=self.tend, dt=self.dt + self.pulse, tstart=tstart, tend=tend, dt=dt ) self.reader(diagnostics) self.equilibrium = self.reader.equilibrium From 67bd35c4e559b90031ecfff836f634e9f48519f6 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 17 Aug 2023 16:02:04 +0100 Subject: [PATCH 153/160] precommit types --- indica/workflows/bayes_workflow_example.py | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/indica/workflows/bayes_workflow_example.py b/indica/workflows/bayes_workflow_example.py index 53809312..bce9272c 100644 --- a/indica/workflows/bayes_workflow_example.py +++ b/indica/workflows/bayes_workflow_example.py @@ -2,6 +2,7 @@ import flatdict import numpy as np from scipy.stats import loguniform +from typing import Optional from indica.bayesmodels import BayesModels from indica.bayesmodels import get_uniform @@ -105,9 +106,9 @@ class BayesWorkflowExample(AbstractBayesWorkflow): def __init__( self, pulse: int = None, - diagnostics: list = None, - param_names: list = None, - opt_quantity: list = None, + diagnostics: Optional[list] = None, + param_names: Optional[list] = None, + opt_quantity: Optional[list] = None, priors: dict = None, profile_params: dict = None, phantoms: bool = False, @@ -144,15 +145,18 @@ def __init__( # TODO: Add some abstraction here if pulse is None: print("Running in test mode") + example_transforms = { + "xrcs": helike_transform_example(1), + "smmh1": smmh1_transform_example(1), + "cxff_pi": pi_transform_example(5), + } self.read_test_data( - { - "xrcs": helike_transform_example(1), - "smmh1": smmh1_transform_example(1), - "cxff_pi": pi_transform_example(5), - } + example_transforms, tstart=self.tstart, tend=self.tend, dt=self.dt ) else: - self.read_data(self.diagnostics) + self.read_data( + self.diagnostics, tstart=self.tstart, tend=self.tend, dt=self.dt + ) self.setup_models(self.diagnostics) def setup_plasma( From d1afc8720a7be7ff60de1607a93d6b0ff64f3dce Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 17 Aug 2023 16:10:26 +0100 Subject: [PATCH 154/160] precommit types --- indica/workflows/bayes_workflow_example.py | 28 ++++++++++++---------- 1 file changed, 15 insertions(+), 13 deletions(-) diff --git a/indica/workflows/bayes_workflow_example.py b/indica/workflows/bayes_workflow_example.py index bce9272c..5d153a01 100644 --- a/indica/workflows/bayes_workflow_example.py +++ b/indica/workflows/bayes_workflow_example.py @@ -1,8 +1,9 @@ +from typing import Optional + import emcee import flatdict import numpy as np from scipy.stats import loguniform -from typing import Optional from indica.bayesmodels import BayesModels from indica.bayesmodels import get_uniform @@ -105,12 +106,12 @@ class BayesWorkflowExample(AbstractBayesWorkflow): def __init__( self, + diagnostics: list, + param_names: list, + opt_quantity: list, + priors: dict, + profile_params: dict, pulse: int = None, - diagnostics: Optional[list] = None, - param_names: Optional[list] = None, - opt_quantity: Optional[list] = None, - priors: dict = None, - profile_params: dict = None, phantoms: bool = False, tstart=0.02, tend=0.10, @@ -414,18 +415,19 @@ def __call__( if __name__ == "__main__": run = BayesWorkflowExample( - pulse=None, - phantoms=True, - diagnostics=["xrcs", "efit", "smmh1", "cxff_pi"], - opt_quantity=[ + ["xrcs", "efit", "smmh1", "cxff_pi"], + [ "xrcs.spectra", "cxff_pi.ti", "efit.wp", "smmh1.ne", ], - param_names=OPTIMISED_PARAMS, - profile_params=DEFAULT_PROFILE_PARAMS, - priors=DEFAULT_PRIORS, + OPTIMISED_PARAMS, + DEFAULT_PROFILE_PARAMS, + DEFAULT_PRIORS, + + pulse=None, + phantoms=True, tstart=0.02, tend=0.10, dt=0.005, From 1f38f27fe762da26a978326376be13fcf40faffa Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 17 Aug 2023 16:23:04 +0100 Subject: [PATCH 155/160] precommit types --- indica/workflows/abstract_bayes_workflow.py | 12 +++------ indica/workflows/bayes_workflow_example.py | 29 +++++++++------------ 2 files changed, 16 insertions(+), 25 deletions(-) diff --git a/indica/workflows/abstract_bayes_workflow.py b/indica/workflows/abstract_bayes_workflow.py index 3221e186..485bc768 100644 --- a/indica/workflows/abstract_bayes_workflow.py +++ b/indica/workflows/abstract_bayes_workflow.py @@ -45,9 +45,7 @@ def read_test_data( self.data: dict = {} def read_data(self, diagnostics: list, tstart=None, tend=None, dt=None): - self.reader = ReadST40( - self.pulse, tstart=tstart, tend=tend, dt=dt - ) + self.reader = ReadST40(self.pulse, tstart=tstart, tend=tend, dt=dt) self.reader(diagnostics) self.equilibrium = self.reader.equilibrium self.transforms = self.reader.transforms @@ -467,15 +465,11 @@ def sample_with_autocorr(sampler, start_points, iterations, n_params, auto_sampl if sampler.iteration % auto_sample: continue new_tau = sampler.get_autocorr_time(tol=0) - autocorr[ - sampler.iteration - 1, - ] = new_tau + autocorr[sampler.iteration - 1,] = new_tau converged = np.all(new_tau * 50 < sampler.iteration) converged &= np.all(np.abs(old_tau - new_tau) / new_tau < 0.01) if converged: break old_tau = new_tau - autocorr = autocorr[ - : sampler.iteration, - ] + autocorr = autocorr[: sampler.iteration,] return autocorr diff --git a/indica/workflows/bayes_workflow_example.py b/indica/workflows/bayes_workflow_example.py index 5d153a01..27a025dc 100644 --- a/indica/workflows/bayes_workflow_example.py +++ b/indica/workflows/bayes_workflow_example.py @@ -1,5 +1,3 @@ -from typing import Optional - import emcee import flatdict import numpy as np @@ -218,8 +216,9 @@ def setup_models(self, diagnostics: list): # passes=2, # ) los_transform.set_equilibrium(self.equilibrium) - model = Interferometry(name=diag) - model.set_los_transform(los_transform) + smmh = Interferometry(name=diag) + smmh.set_los_transform(los_transform) + self.models[diag] = smmh elif diag == "xrcs": los_transform = self.transforms[diag] @@ -229,24 +228,27 @@ def setup_models(self, diagnostics: list): if diag in self.data.keys(): window = self.data[diag]["spectra"].wavelength.values - model = HelikeSpectrometer( + xrcs = HelikeSpectrometer( name="xrcs", window_masks=[slice(0.394, 0.396)], window=window, ) - model.set_los_transform(los_transform) + xrcs.set_los_transform(los_transform) + self.models[diag] = xrcs elif diag == "efit": - model = EquilibriumReconstruction(name="efit") + efit = EquilibriumReconstruction(name="efit") + self.models[diag] = efit elif diag == "cxff_pi": transform = self.transforms[diag] transform.set_equilibrium(self.equilibrium) - model = ChargeExchange(name=diag, element="ar") - model.set_transect_transform(transform) + cxrs = ChargeExchange(name=diag, element="ar") + cxrs.set_transect_transform(transform) + self.models[diag] = cxrs + else: raise ValueError(f"{diag} not found in setup_models") - self.models[diag] = model def setup_opt_data(self, phantoms=False, **kwargs): if not hasattr(self, "plasma"): @@ -294,11 +296,7 @@ def _phantom_data(self, noise=False, noise_factor=0.1, **kwargs): ].max().values * np.random.normal(0, noise_factor, None) opt_data["xrcs.spectra"] = opt_data["xrcs.spectra"] + np.random.normal( 0, - np.sqrt( - opt_data["xrcs.spectra"].values[ - 0, - ] - ), + np.sqrt(opt_data["xrcs.spectra"].values[0,]), opt_data["xrcs.spectra"].shape[1], ) opt_data["cxff_pi.ti"] = opt_data["cxff_pi.ti"] + opt_data[ @@ -425,7 +423,6 @@ def __call__( OPTIMISED_PARAMS, DEFAULT_PROFILE_PARAMS, DEFAULT_PRIORS, - pulse=None, phantoms=True, tstart=0.02, From c555f6b6900ae5dca3e448dcf785e8b1cb5ada40 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 17 Aug 2023 16:27:54 +0100 Subject: [PATCH 156/160] precommit types --- indica/workflows/bayes_workflow_example.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/indica/workflows/bayes_workflow_example.py b/indica/workflows/bayes_workflow_example.py index 27a025dc..a73b0cef 100644 --- a/indica/workflows/bayes_workflow_example.py +++ b/indica/workflows/bayes_workflow_example.py @@ -234,7 +234,7 @@ def setup_models(self, diagnostics: list): window=window, ) xrcs.set_los_transform(los_transform) - self.models[diag] = xrcs + # self.models[diag] = xrcs elif diag == "efit": efit = EquilibriumReconstruction(name="efit") From b5d584010984158a27a43d004510a94d8f6dd158 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 17 Aug 2023 16:37:13 +0100 Subject: [PATCH 157/160] precommit types --- indica/workflows/bayes_workflow_example.py | 21 +++++++++------------ 1 file changed, 9 insertions(+), 12 deletions(-) diff --git a/indica/workflows/bayes_workflow_example.py b/indica/workflows/bayes_workflow_example.py index a73b0cef..8e511f5e 100644 --- a/indica/workflows/bayes_workflow_example.py +++ b/indica/workflows/bayes_workflow_example.py @@ -2,6 +2,7 @@ import flatdict import numpy as np from scipy.stats import loguniform +from typing import Dict, Any from indica.bayesmodels import BayesModels from indica.bayesmodels import get_uniform @@ -196,7 +197,7 @@ def setup_plasma( self.save_phantom_profiles() def setup_models(self, diagnostics: list): - self.models = {} + self.models: Dict[str, Any] = {} for diag in diagnostics: if diag == "smmh1": los_transform = self.transforms[diag] @@ -216,9 +217,8 @@ def setup_models(self, diagnostics: list): # passes=2, # ) los_transform.set_equilibrium(self.equilibrium) - smmh = Interferometry(name=diag) - smmh.set_los_transform(los_transform) - self.models[diag] = smmh + model = Interferometry(name=diag) + model.set_los_transform(los_transform) elif diag == "xrcs": los_transform = self.transforms[diag] @@ -228,27 +228,24 @@ def setup_models(self, diagnostics: list): if diag in self.data.keys(): window = self.data[diag]["spectra"].wavelength.values - xrcs = HelikeSpectrometer( + model = HelikeSpectrometer( name="xrcs", window_masks=[slice(0.394, 0.396)], window=window, ) - xrcs.set_los_transform(los_transform) - # self.models[diag] = xrcs + model.set_los_transform(los_transform) elif diag == "efit": efit = EquilibriumReconstruction(name="efit") - self.models[diag] = efit elif diag == "cxff_pi": transform = self.transforms[diag] transform.set_equilibrium(self.equilibrium) - cxrs = ChargeExchange(name=diag, element="ar") - cxrs.set_transect_transform(transform) - self.models[diag] = cxrs - + model = ChargeExchange(name=diag, element="ar") + model.set_transect_transform(transform) else: raise ValueError(f"{diag} not found in setup_models") + self.models[diag] = model def setup_opt_data(self, phantoms=False, **kwargs): if not hasattr(self, "plasma"): From 9d12492cac32640d15b3277e04b3c74bc249af38 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 17 Aug 2023 16:44:27 +0100 Subject: [PATCH 158/160] precommit types --- indica/workflows/bayes_workflow_example.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/indica/workflows/bayes_workflow_example.py b/indica/workflows/bayes_workflow_example.py index 8e511f5e..ed1a15eb 100644 --- a/indica/workflows/bayes_workflow_example.py +++ b/indica/workflows/bayes_workflow_example.py @@ -198,6 +198,7 @@ def setup_plasma( def setup_models(self, diagnostics: list): self.models: Dict[str, Any] = {} + model: Any = None for diag in diagnostics: if diag == "smmh1": los_transform = self.transforms[diag] @@ -236,7 +237,7 @@ def setup_models(self, diagnostics: list): model.set_los_transform(los_transform) elif diag == "efit": - efit = EquilibriumReconstruction(name="efit") + model = EquilibriumReconstruction(name="efit") elif diag == "cxff_pi": transform = self.transforms[diag] From 15c4aacac7ed989b6ef030d4cc347743b0d98509 Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 17 Aug 2023 16:48:29 +0100 Subject: [PATCH 159/160] precommit types --- indica/workflows/abstract_bayes_workflow.py | 8 ++++++-- indica/workflows/bayes_workflow_example.py | 10 ++++++++-- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/indica/workflows/abstract_bayes_workflow.py b/indica/workflows/abstract_bayes_workflow.py index 485bc768..4839a2b5 100644 --- a/indica/workflows/abstract_bayes_workflow.py +++ b/indica/workflows/abstract_bayes_workflow.py @@ -465,11 +465,15 @@ def sample_with_autocorr(sampler, start_points, iterations, n_params, auto_sampl if sampler.iteration % auto_sample: continue new_tau = sampler.get_autocorr_time(tol=0) - autocorr[sampler.iteration - 1,] = new_tau + autocorr[ + sampler.iteration - 1, + ] = new_tau converged = np.all(new_tau * 50 < sampler.iteration) converged &= np.all(np.abs(old_tau - new_tau) / new_tau < 0.01) if converged: break old_tau = new_tau - autocorr = autocorr[: sampler.iteration,] + autocorr = autocorr[ + : sampler.iteration, + ] return autocorr diff --git a/indica/workflows/bayes_workflow_example.py b/indica/workflows/bayes_workflow_example.py index ed1a15eb..c52bbf23 100644 --- a/indica/workflows/bayes_workflow_example.py +++ b/indica/workflows/bayes_workflow_example.py @@ -1,8 +1,10 @@ +from typing import Any +from typing import Dict + import emcee import flatdict import numpy as np from scipy.stats import loguniform -from typing import Dict, Any from indica.bayesmodels import BayesModels from indica.bayesmodels import get_uniform @@ -294,7 +296,11 @@ def _phantom_data(self, noise=False, noise_factor=0.1, **kwargs): ].max().values * np.random.normal(0, noise_factor, None) opt_data["xrcs.spectra"] = opt_data["xrcs.spectra"] + np.random.normal( 0, - np.sqrt(opt_data["xrcs.spectra"].values[0,]), + np.sqrt( + opt_data["xrcs.spectra"].values[ + 0, + ] + ), opt_data["xrcs.spectra"].shape[1], ) opt_data["cxff_pi.ti"] = opt_data["cxff_pi.ti"] + opt_data[ From 4362cd5bf7745f8750c93762092c3735415aafde Mon Sep 17 00:00:00 2001 From: michael gemmell Date: Thu, 17 Aug 2023 17:05:07 +0100 Subject: [PATCH 160/160] moving kwargs/args around --- indica/workflows/bayes_workflow_example.py | 15 +++++---------- .../unit/workflows/test_bayes_workflow_example.py | 6 +++--- 2 files changed, 8 insertions(+), 13 deletions(-) diff --git a/indica/workflows/bayes_workflow_example.py b/indica/workflows/bayes_workflow_example.py index c52bbf23..cb81a09d 100644 --- a/indica/workflows/bayes_workflow_example.py +++ b/indica/workflows/bayes_workflow_example.py @@ -417,17 +417,12 @@ def __call__( if __name__ == "__main__": run = BayesWorkflowExample( - ["xrcs", "efit", "smmh1", "cxff_pi"], - [ - "xrcs.spectra", - "cxff_pi.ti", - "efit.wp", - "smmh1.ne", - ], - OPTIMISED_PARAMS, - DEFAULT_PROFILE_PARAMS, - DEFAULT_PRIORS, pulse=None, + diagnostics=["xrcs", "efit", "smmh1", "cxff_pi"], + param_names=OPTIMISED_PARAMS, + opt_quantity=OPTIMISED_QUANTITY, + priors=DEFAULT_PRIORS, + profile_params=DEFAULT_PROFILE_PARAMS, phantoms=True, tstart=0.02, tend=0.10, diff --git a/tests/unit/workflows/test_bayes_workflow_example.py b/tests/unit/workflows/test_bayes_workflow_example.py index d725eb47..383f19c8 100644 --- a/tests/unit/workflows/test_bayes_workflow_example.py +++ b/tests/unit/workflows/test_bayes_workflow_example.py @@ -64,7 +64,7 @@ def teardown_method(self): self.workflow = None def test_workflow_initializes(self): - attributes_to_check = ["data", "reader", "models", "equilibrium"] + attributes_to_check = ["data", "models", "equilibrium"] for attribute in attributes_to_check: if not hasattr(self.workflow, attribute): raise ValueError(f"missing {attribute} in workflow object") @@ -72,11 +72,11 @@ def test_workflow_initializes(self): def test_init_phantoms_false_with_example_plasma(self): with pytest.raises(ValueError): - BayesWorkflowExample(dict(self.init_settings, **{"phantoms": False})) + BayesWorkflowExample(**dict(self.init_settings, **{"phantoms": False})) def test_init_not_including_all_required_inputs(self): with pytest.raises(ValueError): - BayesWorkflowExample(dict(self.init_settings, **{"param_names": None})) + BayesWorkflowExample(**dict(self.init_settings, **{"param_names": None})) # def test_reader_has_read_all_diagnostic_data(self): # assert all(diag_name in self.workflow.reader.keys()