From 3317b2802d9e50e2a3646eeabb60e4d87cafbc56 Mon Sep 17 00:00:00 2001 From: michael-gemmell <87419247+michael-gemmell@users.noreply.github.com> Date: Thu, 17 Aug 2023 17:12:48 +0100 Subject: [PATCH] Michaelgemmell/bayes workflow (#262) * initial commit * magnetic recon tests * use error attribute if it is given * Added poisson noise and adjusted calibration factor * changed name magnetic_recon -> equilibrium_reconstruction * diagnostic data with channel dim are treated as individual measurements * added function of example LOS * refactoring * refactoring * refactoring * initial commit * removed obsolete attributes * Setting up testflow * more tests * more test stubs * small fix * cosmetic changes * changed marker style * small fix * adjusting alpha * black reformatting * fixing formatting * flake8 fix * flake8 fix * boolean indexing -> xarray.where in call * fixing tests * removing import * precommit fixes * stashing for quick result * added calc_impurity for generating profiles from concentrations to update profiles method * added conditional handling of diagnostics in data pipeline * moved calc_impurity to plasma method * tweaking calibration factor * tweaking calibration factor * tweaking calibration factor * changed burn_in to burn_in_fraction * added violinplots * removed comment * included dictionary logic for fitting impurities scaled by concentrations and electron density * changed name cxff_pi.ti_0d to cxff_pi.ti0 * new option for initialising at maximum-likelihood estimate * added xlim for violin plot * tweaking calibration * fixed neutral density units * added set plot functions * added center mass sampling and Ti/Te y0_ref options * plot formatting changed * black formatting * deleted example bayes_opt * moved sample with autocorr function to bayes_dev_workflow * renamed bayesworkflow * stashing * multiplot formatting * generalised writing kinetic profiles to blobs * function autocorr plot added * ion density added to kinetic profiles saved * inital commit of abstract workflow and example of usage * doc strings * renaming workflow * Ti_ref as default * abstract methods added * Changed to new MDSPlus structure * Adding default methods * moved sampling to it's own function in abstract class * moved sampling to it's own function in abstract class * renamed abstract class to AbstractBayesWorkflow * comments * added high density sampling as method * added attributes to __init__ * error handling __init__ * error handling __init__ * moved read data to abstract class * initial commit * kwargs added * xrcs wavelength units corrected * background to int * stash * renamed impurities * stashing * stashing * stashing * fixing key names * updating and testing example * added background as attribute of helike model * fixed phantom methods * rearranged methods * moved background to call kwarg * params and kwargs now are given as model_{var} and model prefix is removed before model call * adding _phantom_data and _exp_data abstract methods * renaming to BayesWorkflowExample * initial commit * renamed phantom_params to profile_params * renamed bayesopt -> bayesmodel * renamed bayesopt -> bayesmodel * moved start point sampling to its own method _sample_start_points * fixed variable name * fixed bug with start_points being overwritten * removed workflow_dev * moving equilibrium to workflow object * when reading raw data save transforms to their own attribute * transforms saved to workflow class * fixed name * renamed example transform * made example los function * minor name fix * made read_test_data function * made example los function * minor type * added reader to read_test_data method * fixed copying workflow object states * updated example to work with assign_profiles * can now run with pulse = None and synthetic transforms/equilibrium * formatting * refactoring names * deleted old bayes_models tests * reformatted _build_bckc for readability * moved percentage error to class attribute * fixing init percent error * black formatting * renamed kin_prof -> plasma_profiles * refactored window handling * replaced doppler_broaden with physics.ev_doppler * removed methods from __init__ * stashing * setup_plasma now takes kwargs * adding plasma to models now happens in setup_opt_data * workflow broken up into methods * removed redundant kwargs * fixed non_plasma call option * fixed plasma_initialisation * fixing violin plotting * units for xrcs.spectra.wavelength fixed * nchannels added to example_los * adding print message for fake data reading * adjusting example los * bckc method not printing spectra/fit not available everytime * adjusting priors * moved kwargs to sample function * nsamples kwarg added to sample_from_high_density_region * mocked bda_tree module * fixed import * fixed import * black formatting * precommit --- indica/bayesmodels.py | 148 ++++-- indica/models/charge_exchange.py | 36 +- indica/models/helike_spectroscopy.py | 87 ++-- indica/models/interferometry.py | 34 +- indica/models/plasma.py | 48 +- indica/readers/read_st40.py | 4 +- indica/workflows/abstract_bayes_workflow.py | 479 ++++++++++++++++++ indica/workflows/bayes_plots.py | 450 ++++++++++++++++ indica/workflows/bayes_workflow.py | 317 ------------ indica/workflows/bayes_workflow_example.py | 445 ++++++++++++++++ indica/workflows/example_bayes_opt.py | 182 ------- indica/workflows/load_modelling_plasma.py | 4 +- indica/workflows/profile_workflows.py | 1 - indica/workflows/run_tomo_1d.py | 2 - indica/workflows/zeff_profile.py | 3 - indica/writers/bda_tree.py | 269 ++++++++++ tests/conftest.py | 1 + .../models/test_diagnostic_models.py | 0 .../models/test_helike_spectroscopy.py | 81 +-- tests/integration/models/test_plasma.py | 6 + tests/integration/test_bayesmodels.py | 74 --- tests/unit/models/__init__.py | 0 tests/unit/models/test_plasma.py | 82 ++- tests/unit/workflows/__init__.py | 0 .../workflows/test_bayes_workflow_example.py | 160 ++++++ 25 files changed, 2141 insertions(+), 772 deletions(-) create mode 100644 indica/workflows/abstract_bayes_workflow.py create mode 100644 indica/workflows/bayes_plots.py delete mode 100644 indica/workflows/bayes_workflow.py create mode 100644 indica/workflows/bayes_workflow_example.py delete mode 100644 indica/workflows/example_bayes_opt.py create mode 100644 indica/writers/bda_tree.py rename tests/{unit => integration}/models/test_diagnostic_models.py (100%) create mode 100644 tests/integration/models/test_plasma.py delete mode 100644 tests/integration/test_bayesmodels.py create mode 100644 tests/unit/models/__init__.py create mode 100644 tests/unit/workflows/__init__.py create mode 100644 tests/unit/workflows/test_bayes_workflow_example.py diff --git a/indica/bayesmodels.py b/indica/bayesmodels.py index a6d92921..a05e8c77 100644 --- a/indica/bayesmodels.py +++ b/indica/bayesmodels.py @@ -1,15 +1,25 @@ from copy import deepcopy import warnings -import flatdict import numpy as np from scipy.stats import uniform np.seterr(all="ignore") - warnings.simplefilter("ignore", category=FutureWarning) +PROFILES = [ + "electron_temperature", + "electron_density", + "ion_temperature", + "ion_density", + "impurity_density", + "fast_density", + "neutral_density", + "zeff", +] + + def gaussian(x, mean, sigma): return 1 / (sigma * np.sqrt(2 * np.pi)) * np.exp(-1 / 2 * ((x - mean) / sigma) ** 2) @@ -44,46 +54,85 @@ def __init__( quant_to_optimise: list = [], priors: dict = {}, diagnostic_models: list = [], + percent_error: float = 0.10, ): 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 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 - # Params is a dictionary which is updated by optimiser, - # kwargs is constant i.e. settings for models + def _build_bckc(self, params, **kwargs): + """ + Parameters + ---------- + params - dictionary which is updated by optimiser + kwargs - passed to model i.e. settings + + Returns + ------- + bckc of results + """ self.bckc: dict = {} for model in self.diagnostic_models: - self.bckc = dict( - self.bckc, **{model.name: {**model(**{**params, **kwargs})}} - ) - self.bckc = flatdict.FlatDict(self.bckc, delimiter=".") + # 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 + _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): 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("float64") + # 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) - .values.astype("float64") + self.data[key].sel(t=self.plasma.time_to_calculate).astype("float128") ) - _ln_likelihood = np.log(gaussian(model_data, exp_data, exp_data * 0.10)) - ln_likelihood += np.nanmean(_ln_likelihood) + 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 + ) + + _ln_likelihood = np.log(gaussian(model_data, exp_data, exp_error)) + # 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): @@ -144,6 +193,34 @@ 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, 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 @@ -164,7 +241,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) @@ -172,23 +249,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 = { - "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 - ), - "zeff": self.plasma.zeff.sum("element").sel( - t=self.plasma.time_to_calculate - ), - # TODO: add Nh - } - blob = deepcopy({**self.bckc, **kin_profs}) + plasma_profiles = {} + for profile_key in PROFILES: + if hasattr(self.plasma, profile_key): + 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, **plasma_profiles}) return ln_posterior, blob diff --git a/indica/models/charge_exchange.py b/indica/models/charge_exchange.py index 61619348..5551d8a6 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 @@ -43,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 @@ -115,31 +120,34 @@ 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, diagnostic_name: str = "cxrs", plasma=None, plot=False, ): - # TODO: LOS sometimes crossing bad EFIT reconstruction if plasma is None: 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, diff --git a/indica/models/helike_spectroscopy.py b/indica/models/helike_spectroscopy.py index 34520aae..90ac5c4e 100644 --- a/indica/models/helike_spectroscopy.py +++ b/indica/models/helike_spectroscopy.py @@ -24,7 +24,7 @@ } -class Helike_spectroscopy(DiagnosticModel): +class HelikeSpectrometer(DiagnosticModel): """ Data and methods to model XRCS spectrometer measurements @@ -35,13 +35,14 @@ def __init__( self, name: str, instrument_method="get_helike_spectroscopy", + etendue: float = 1.0, + calibration: float = 8.0e-20, element: str = "ar", window_len: int = 1030, - window_lim: list = [0.394, 0.401], - window_masks: list = [], - etendue: float = 1.0, - calibration: float = 1.0e-27, - line_labels: list = ["w", "k", "n3", "n345", "z", "qra"], + window_lim=None, + window: np.array = None, + window_masks=None, + line_labels=None, ): """ Read all atomic data and initialise objects @@ -50,25 +51,32 @@ def __init__( ---------- name 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 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 window_masks: - for mslice in window_masks: + if self.window_masks: + for mslice in self.window_masks: mask[(window > mslice.start) & (window < mslice.stop)] = 1 else: mask[:] = 1 @@ -377,6 +385,7 @@ def __call__( t: LabeledArray = None, calc_rho: bool = False, moment_analysis: bool = False, + background: int = None, **kwargs, ): """ @@ -404,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( @@ -444,17 +453,43 @@ def __call__( self.quantities: dict = AVAILABLE_QUANTITIES[self.instrument_method] # TODO: check that inputs have compatible dimensions/coordinates - self._calculate_intensity() self._make_spectra() if moment_analysis: self._moment_analysis() + if background is not None: + self.measured_spectra = self.measured_spectra + background + 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 + los_end[:, 1] = 0.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 + + 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="xrcs", + machine_dimensions=((0.15, 0.95), (-0.7, 0.7)), + passes=1, + ) + return los_transform + + def example_run( pulse: int = None, plasma=None, plot=False, moment_analysis: bool = False, **kwargs ): @@ -466,28 +501,9 @@ def example_run( plasma.time_to_calculate = plasma.t[5] # Create new diagnostic diagnostic_name = "xrcs" - 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], - origin[:, 1], - origin[:, 2], - direction[:, 0], - direction[:, 1], - direction[:, 2], - name=diagnostic_name, - machine_dimensions=plasma.machine_dimensions, - passes=1, - ) + los_transform = helike_transform_example(3) los_transform.set_equilibrium(plasma.equilibrium) - model = Helike_spectroscopy( + model = HelikeSpectrometer( diagnostic_name, window_masks=[], ) @@ -518,8 +534,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() diff --git a/indica/models/interferometry.py b/indica/models/interferometry.py index 3ffac6e1..5a7cf49b 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] @@ -90,19 +89,13 @@ 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" - 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 - model = Interferometry( - diagnostic_name, - ) los_transform = LineOfSightTransform( origin[:, 0], origin[:, 1], @@ -110,10 +103,23 @@ 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(3) los_transform.set_equilibrium(plasma.equilibrium) model.set_los_transform(los_transform) model.set_plasma(plasma) diff --git a/indica/models/plasma.py b/indica/models/plasma.py index fd36d785..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 @@ -202,7 +201,6 @@ def initialize_variables(self): (to be used e.g. in optimisation workflows) """ - # Dictionary keeping track of deta use for optimisations self.optimisation: dict = {} self.forward_models: dict = {} self.power_loss_sxr: dict = {} @@ -304,7 +302,7 @@ def initialize_variables(self): 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") @@ -495,18 +493,18 @@ def assign_profiles( def update_profiles( self, parameters: dict, + ): + """ + 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", - ], - ): - """ - Update plasma profiles with profile parameters i.e. - {"Ne_prof.y0":1e19} -> Ne_prof.y0 - """ + ] for param, value in parameters.items(): _prefix = [pref for pref in profile_prefixs if pref in param] if _prefix: @@ -857,22 +855,17 @@ 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): """ @@ -1272,6 +1265,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) @@ -1287,18 +1281,18 @@ 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.assign_profiles("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.assign_profiles("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.assign_profiles("toroidal_rotation", t=t) plasma.Ne_prof.peaking = ne_peaking[i] - plasma.assign_profiles(profile="electron_density", t=t) + plasma.assign_profiles("electron_density", t=t) plasma.Nimp_prof.peaking = nimp_peaking[i] plasma.Nimp_prof.y0 = nimp_y0[i] diff --git a/indica/readers/read_st40.py b/indica/readers/read_st40.py index fcbb8fd3..62cdaf62 100644 --- a/indica/readers/read_st40.py +++ b/indica/readers/read_st40.py @@ -94,6 +94,7 @@ def __init__( self.equilibrium: Equilibrium self.raw_data: dict = {} self.binned_data: dict = {} + self.transforms: dict = {} def reset_data(self): self.raw_data = {} @@ -121,6 +122,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 @@ -173,7 +175,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!" @@ -308,7 +309,6 @@ def __call__( debug: bool = False, ): self.debug = debug - 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 new file mode 100644 index 00000000..4839a2b5 --- /dev/null +++ b/indica/workflows/abstract_bayes_workflow.py @@ -0,0 +1,479 @@ +from abc import ABC +from abc import abstractmethod +from pathlib import Path +import pickle + +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): + @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_models(self.diagnostics) + + 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( + tstart, + tend, + dt, + ) + self.transforms = diagnostic_transforms + 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(diagnostics) + self.equilibrium = self.reader.equilibrium + self.transforms = self.reader.transforms + 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) + self.save_phantom_profiles() + + @abstractmethod + def setup_models(self, diagnostics: list): + """ + Initialising models normally requires data to be read so transforms can be set + + """ + self.models: dict = {} + + @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): + """ + 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: + self.opt_data = self._exp_data() + + @abstractmethod + def setup_optimiser(self, model_kwargs): + """ + Initialise and provide settings for optimiser + """ + 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): + """ + + 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 + + result["METADATA"] = { + "GITCOMMIT": "PLACEHOLDER", + "USER": "PLACEHOLDER", + "EQUIL": "PLACEHOLDER", + } + + result["INPUT"] = { + "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, + "MAIN_ION": self.plasma.main_ion, + } + result["INPUT"]["WORKFLOW"] = { + diag_name.upper(): { + "PULSE": self.pulse, # Change this if different pulses used + "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"].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["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"), + "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"), + } + 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"].sel(element=self.plasma.main_ion), + "TE": self.blobs["electron_temperature"], + "TI": self.blobs["ion_temperature"].sel(element=self.plasma.main_ion), + "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["OPTIMISATION"] = { + "ACCEPT_FRAC": self.accept_frac, + "PRIOR_SAMPLE": self.prior_sample, + "POST_SAMPLE": self.post_sample, + "AUTOCORR": self.autocorr, + } + + 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"), + } + 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, iterations, burn_frac): + """ + 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.burn_frac = burn_frac + self.iterations = iterations + + self.autocorr = sample_with_autocorr( + self.sampler, + self.start_points, + iterations, + self.param_names.__len__(), + auto_sample=10, + ) + 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__()) + + 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.bayesmodel.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) + + 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 + + +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 diff --git a/indica/workflows/bayes_plots.py b/indica/workflows/bayes_plots.py new file mode 100644 index 00000000..93955202 --- /dev/null +++ b/indica/workflows/bayes_plots.py @@ -0,0 +1,450 @@ +import os +import pickle + +import corner +import flatdict +import matplotlib.pyplot as plt +import numpy as np + +from indica.utilities import set_axis_sci +from indica.utilities import set_plot_rcparams + + +def plot_profile( + profile, + blobkey: str, + figheader="./results/test/", + phantom_profile=None, + logscale=False, + sharefig=False, + filename="", + filetype=".png", + linestyle="--", + color="blue", +): + set_plot_rcparams("profiles") + + plt.fill_between( + profile.rho_poloidal, + profile.quantile(0.16, dim="index"), + profile.quantile(0.84, dim="index"), + label=f"{blobkey}, 68% Confidence", + zorder=3, + color=color, + alpha=0.9, + ) + if blobkey != "NFAST": + plt.fill_between( + profile.rho_poloidal, + profile.quantile(0.025, dim="index"), + profile.quantile(0.975, dim="index"), + label=f"{blobkey}, 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"{blobkey}, Max-Min", + zorder=1, + color="lightgrey", + alpha=0.2, + ) + + 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"{blobkey}, phantom profile", + linestyle=linestyle, + color="black", + zorder=4, + ) + + plt.legend() + 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}{filetype}") + else: + plt.savefig(figheader + f"{blobkey}{filetype}") + plt.close("all") + + +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, +): + if blobkey not in blobs.keys(): + raise ValueError(f"{blobkey} not in blobs") + + plt.figure(figsize=figsize) + blob_data = blobs[blobkey] + dims = tuple(name for name in blob_data.dims if name != "index") + plt.fill_between( + blob_data.__getattr__(dims[0]), + blob_data.quantile(0.16, dim="index"), + blob_data.quantile(0.84, dim="index"), + label=f"{blobkey}, 68% Confidence", + zorder=3, + color="blue", + ) + plt.fill_between( + blob_data.__getattr__(dims[0]), + blob_data.quantile(0.025, dim="index"), + blob_data.quantile(0.975, dim="index"), + label=f"{blobkey}, 95% Confidence", + zorder=2, + color="grey", + ) + plt.fill_between( + blob_data.__getattr__(dims[0]), + blob_data.quantile(0.00, dim="index"), + blob_data.quantile(1.00, dim="index"), + label=f"{blobkey}, Max-Min", + zorder=1, + color="lightgrey", + ) + 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, + ) + set_axis_sci() + plt.ylabel(ylabel) + plt.xlabel(xlabel) + plt.xlim(xlim) + plt.legend() + plt.savefig(figheader + filename) + plt.close() + + +def violinplot( + data: dict, + key: str, + diag_data: dict, + filename: str, + ylabel="[a.u.]", + figheader="./results/test/", + **kwargs, +): + set_plot_rcparams("multi") + 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, + ) + 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}") + + 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()): + 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="-.", + label=f"{key} data", + ) + axs[0].set_ylabel("pdf ()") + + plt.savefig(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") + plt.legend(param_names) + 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/", + filetype=".png", + **kwargs, +): + # 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 + ) + + key = "CXFF_PI.TI0" + if key in model_data.keys(): + violinplot( + model_data, + key, + diag_data, + f"{key.replace('.', '_')}" + filetype, + figheader=figheader, + ylabel="Temperature [eV]", + ) + + key = "EFIT.WP" + if key in model_data.keys(): + violinplot( + model_data, + key, + diag_data, + f"{key.replace('.', '_')}" + filetype, + figheader=figheader, + ylabel="Energy [J]", + ) + key = "SMMH1.NE" + if key in model_data.keys(): + violinplot( + 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 model_data.keys(): + violinplot( + model_data, + key, + diag_data, + f"{key.replace('.', '_')}" + filetype, + figheader=figheader, + ylabel="Temperature [eV]", + ) + key = "XRCS.TI_W" + if key in model_data.keys(): + violinplot( + model_data, + key, + diag_data, + f"{key.replace('.', '_')}" + filetype, + figheader=figheader, + ylabel="Temperature [eV]", + ) + + set_plot_rcparams("multi") + key = "XRCS.SPECTRA" + if key in model_data.keys(): + _plot_1d( + model_data, + key, + diag_data, + f"{key.replace('.', '_')}" + filetype, + figheader=figheader, + ylabel="Intensity [a.u.]", + xlabel="Wavelength [nm]", + xlim=(0.394, 0.401), + figsize=(6, 4), + ) + key = "CXFF_PI.TI" + if key in model_data.keys(): + _plot_1d( + model_data, + key, + diag_data, + f"{key.replace('.', '_')}" + filetype, + figheader=figheader, + ylabel="Temperature [eV]", + xlabel="Channel", + ) + + key = "TE" + plot_profile( + profiles[key], + key, + figheader=figheader, + filetype=filetype, + phantom_profile=phantom_profiles, + sharefig=True, + color="blue", + linestyle="dashdot", + ) + key = "TI" + plot_profile( + profiles[key], + key, + figheader=figheader, + filename="temperature", + filetype=filetype, + phantom_profile=phantom_profiles, + color="red", + linestyle="dotted", + ) + + key = "NE" + plot_profile( + profiles[key], + key, + figheader=figheader, + filetype=filetype, + phantom_profile=phantom_profiles, + color="blue", + sharefig=True, + ) + key = "NI" + plot_profile( + profiles[key], + key, + figheader=figheader, + filetype=filetype, + phantom_profile=phantom_profiles, + sharefig=True, + color="red", + ) + key = "NFAST" + plot_profile( + profiles[key], + key, + figheader=figheader, + filename="densities", + filetype=filetype, + phantom_profile=phantom_profiles, + color="green", + ) + + key = "NIZ1" + plot_profile( + profiles[key], + key, + figheader=figheader, + filename="ar density", + filetype=filetype, + phantom_profile=phantom_profiles, + color="red", + ) + + key = "NIZ2" + plot_profile( + profiles[key], + key, + figheader=figheader, + filename="c density", + filetype=filetype, + phantom_profile=phantom_profiles, + color="red", + ) + + key = "NNEUTR" + plot_profile( + profiles[key], + key, + filename="neutral density", + figheader=figheader, + filetype=filetype, + phantom_profile=phantom_profiles, + logscale=True, + ) + + corner.corner(post_sample, labels=param_names) + plt.savefig(figheader + "posterior" + filetype) + + corner.corner(prior_sample, labels=param_names) + plt.savefig(figheader + "prior" + filetype) + plt.close("all") + + +if __name__ == "__main__": + filehead = "./results/test/" + with open(filehead + "results.pkl", "rb") as handle: + results = pickle.load(handle) + plot_bayes_result(results, filehead, filetype=".png") diff --git a/indica/workflows/bayes_workflow.py b/indica/workflows/bayes_workflow.py deleted file mode 100644 index 91e50965..00000000 --- a/indica/workflows/bayes_workflow.py +++ /dev/null @@ -1,317 +0,0 @@ -from pathlib import Path -import time - -import corner -import matplotlib.pyplot as plt -import numpy as np - -timestr = time.strftime("%Y%m%d%H%M") - - -def plot_profile( - profile, - blobkey: str, - figheader="./results/test/", - phantom_profile=None, - sharefig=False, - filename="", - linestyle="--", - color="blue", -): - if not plt.get_fignums(): # If no figure is open - plt.figure(figsize=(8, 6)) - - plt.fill_between( - profile.rho_poloidal, - profile.quantile(0.16, dim="index"), - profile.quantile(0.84, dim="index"), - label=f"{blobkey}, 68% Confidence", - zorder=3, - color=color, - alpha=0.9, - ) - plt.fill_between( - profile.rho_poloidal, - profile.quantile(0.025, dim="index"), - profile.quantile(0.975, dim="index"), - label=f"{blobkey}, 95% Confidence", - zorder=2, - color="grey", - alpha=0.7, - ) - plt.fill_between( - profile.rho_poloidal, - profile.quantile(0.00, dim="index"), - profile.quantile(1.00, dim="index"), - label=f"{blobkey}, Max-Min", - zorder=1, - color="lightgrey", - alpha=0.7, - ) - - if phantom_profile is not None: - phantom_profile.plot( - label=f"{blobkey}, phantom profile", - linestyle=linestyle, - color="black", - zorder=4, - ) - - plt.legend() - if sharefig: - return - - if filename: - plt.savefig(figheader + timestr + f"{filename}.png") - else: - plt.savefig(figheader + timestr + f"{blobkey}.png") - plt.close("all") - - -def _plot_0d( - blobs: dict, - blobkey: str, - diag_data: dict, - filename: str, - figheader="./results/test/", - xlabel="samples ()", - ylabel="a.u.", - **kwargs, -): - if blobkey not in blobs.keys(): - raise ValueError(f"{blobkey} not in blobs") - plt.figure() - 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.", - **kwargs, -): - if blobkey not in blobs.keys(): - raise ValueError(f"{blobkey} not in blobs") - - plt.figure() - blob_data = blobs[blobkey] - dims = tuple(name for name in blob_data.dims if name != "index") - plt.fill_between( - blob_data.__getattr__(dims[0]), - blob_data.quantile(0.16, dim="index"), - blob_data.quantile(0.84, dim="index"), - label=f"{blobkey}, 68% Confidence", - zorder=3, - color="blue", - ) - plt.fill_between( - blob_data.__getattr__(dims[0]), - blob_data.quantile(0.025, dim="index"), - blob_data.quantile(0.975, dim="index"), - label=f"{blobkey}, 95% Confidence", - zorder=2, - color="grey", - ) - plt.fill_between( - blob_data.__getattr__(dims[0]), - blob_data.quantile(0.00, dim="index"), - blob_data.quantile(1.00, dim="index"), - label=f"{blobkey}, Max-Min", - 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, - ) - plt.ylabel(ylabel) - plt.xlabel(dims[0]) - plt.legend() - plt.savefig(figheader + 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, -): - Path(figheader).mkdir(parents=True, exist_ok=True) - - 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 + timestr + "average_tau.png") - plt.close() - - key = "pi.brightness" - if key in blobs.keys(): - _plot_1d( - blobs, - key, - diag_data, - f"{timestr}_{key.replace('.', '_')}.png", - figheader=figheader, - ylabel="Intensity (W m^-2)", - ) - key = "efit.wp" - if key in blobs.keys(): - _plot_0d( - blobs, - key, - diag_data, - f"{key.replace('.', '_')}.png", - figheader=figheader, - ylabel="Wp (J)", - ) - key = "smmh1.ne" - if key in blobs.keys(): - _plot_0d( - blobs, - key, - diag_data, - f"{timestr}_{key.replace('.', '_')}.png", - figheader=figheader, - ylabel="ne_int (m^-2)", - ) - key = "xrcs.te_kw" - if key in blobs.keys(): - _plot_0d( - blobs, - key, - diag_data, - f"{key.replace('.', '_')}.png", - figheader=figheader, - ylabel="temperature (eV)", - ) - key = "xrcs.ti_w" - if key in blobs.keys(): - _plot_0d( - blobs, - key, - diag_data, - f"{key.replace('.', '_')}.png", - figheader=figheader, - ylabel="temperature (eV)", - ) - key = "xrcs.spectra" - if key in blobs.keys(): - _plot_1d( - blobs, - key, - diag_data, - f"{key.replace('.', '_')}.png", - figheader=figheader, - ylabel="intensity (A.U.)", - ) - key = "cxrs.ti" - if key in blobs.keys(): - _plot_1d( - blobs, - key, - diag_data, - f"{key.replace('.', '_')}.png", - figheader=figheader, - ylabel="temperature (eV)", - ) - - key = "zeff" - if key in blobs and key in phantom_profiles: - plot_profile( - blobs[key], key, figheader=figheader, phantom_profile=phantom_profiles[key] - ) - plt.ylim(0, 10) - - key = "electron_temperature" - plot_profile( - blobs[key], - key, - figheader=figheader, - phantom_profile=phantom_profiles[key], - sharefig=True, - color="blue", - linestyle="dashdot", - ) - key = "ion_temperature" - element = blobs[key].element[0] - plot_profile( - blobs[key].sel(element=element), - key, - figheader=figheader, - filename="temperature", - phantom_profile=phantom_profiles[key], - color="red", - linestyle="dotted", - ) - key = "electron_density" - plot_profile( - blobs[key], key, figheader=figheader, phantom_profile=phantom_profiles[key] - ) - key = "impurity_density" - impurity = blobs[key].element[0] - plot_profile( - blobs[key].sel(element=impurity), - key, - figheader=figheader, - phantom_profile=phantom_profiles[key], - color="red", - ) - - corner.corner(samples, labels=param_names) - plt.savefig(figheader + timestr + "posterior.png") - - corner.corner(prior_samples, labels=param_names) - plt.savefig(figheader + timestr + "prior.png") - 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 diff --git a/indica/workflows/bayes_workflow_example.py b/indica/workflows/bayes_workflow_example.py new file mode 100644 index 00000000..cb81a09d --- /dev/null +++ b/indica/workflows/bayes_workflow_example.py @@ -0,0 +1,445 @@ +from typing import Any +from typing import Dict + +import emcee +import flatdict +import numpy as np +from scipy.stats import loguniform + +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 = { + "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": 1e17, + "Nimp_prof.y1": 5e15, + "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), + "Nimp_prof.y0": loguniform(1e16, 1e18), + "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, 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), + "Ti_prof.wcenter": get_uniform(0.1, 0.8), + "Ti_prof.peaking": get_uniform(1, 10), +} + +OPTIMISED_PARAMS = [ + # "Ne_prof.y1", + "Ne_prof.y0", + # "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", +] + + +class BayesWorkflowExample(AbstractBayesWorkflow): + def __init__( + self, + diagnostics: list, + param_names: list, + opt_quantity: list, + priors: dict, + profile_params: dict, + pulse: int = None, + phantoms: bool = False, + tstart=0.02, + tend=0.10, + dt=0.005, + ): + self.pulse = pulse + self.diagnostics = diagnostics + self.param_names = param_names + self.opt_quantity = opt_quantity + 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", + "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" + ) + + # 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( + example_transforms, tstart=self.tstart, tend=self.tend, dt=self.dt + ) + else: + self.read_data( + self.diagnostics, tstart=self.tstart, tend=self.tend, dt=self.dt + ) + 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, + ): + 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=main_ion, + impurities=impurities, + impurity_concentration=impurity_concentration, + full_run=False, + n_rad=n_rad, + ) + self.tsample = tsample + self.plasma.time_to_calculate = self.plasma.t[ + 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: Dict[str, Any] = {} + model: Any = None + 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.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.equilibrium) + window = None + if hasattr(self, "data"): + if diag in self.data.keys(): + window = self.data[diag]["spectra"].wavelength.values + + model = HelikeSpectrometer( + name="xrcs", + window_masks=[slice(0.394, 0.396)], + window=window, + ) + model.set_los_transform(los_transform) + + elif diag == "efit": + model = EquilibriumReconstruction(name="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) + 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"): + 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(**kwargs) + else: + self.opt_data = self._exp_data(**kwargs) + + def _phantom_data(self, noise=False, noise_factor=0.1, **kwargs): + opt_data = {} + if "smmh1" in self.diagnostics: + opt_data["smmh1.ne"] = ( + self.models["smmh1"]() + .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"]) + if "cxff_pi" in self.diagnostics: + cxrs_data = ( + self.models["cxff_pi"]() + .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]}) + ) + + if noise: + 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( + 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( + 0, noise_factor, opt_data["cxff_pi.ti"].shape[1] + ) + opt_data["efit.wp"] = opt_data["efit.wp"] + opt_data[ + "efit.wp" + ].max().values * np.random.normal(0, noise_factor, None) + return opt_data + + 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 + ) + 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) + opt_data["xrcs.spectra"]["error"] = np.sqrt( + opt_data["xrcs.spectra"] + background.std(dim="wavelength") ** 2 + ) + + if "cxff_pi" in self.diagnostics: + 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, + nwalkers=50, + burn_frac=0.10, + sample_high_density=False, + ): + self.model_kwargs = model_kwargs + 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, + 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( + nwalkers, + ndim, + log_prob_fn=self.bayesmodel.ln_posterior, + parameter_names=self.param_names, + moves=self.move, + kwargs=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): + if sample_high_density: + start_points = self.bayesmodel.sample_from_high_density_region( + self.param_names, self.sampler, self.nwalkers, nsamples=100 + ) + else: + start_points = self.bayesmodel.sample_from_priors( + self.param_names, size=self.nwalkers + ) + 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, + ): + if mds_write: + # check_analysis_run(self.pulse, self.run) + self.node_structure = create_nodes( + pulse_to_write=pulse_to_write, + run=run, + diagnostic_quantities=self.opt_quantity, + mode="NEW", + ) + + self.run_sampler(iterations=iterations, burn_frac=burn_frac) + self.save_pickle(filepath=filepath) + + if plot: # currently requires result with DataArrays + plot_bayes_result(self.result, filepath) + + self.result = self.dict_of_dataarray_to_numpy(self.result) + if mds_write: + write_nodes(pulse_to_write, self.node_structure, self.result) + + return self.result + + +if __name__ == "__main__": + run = BayesWorkflowExample( + 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, + dt=0.005, + ) + + run.setup_plasma( + tsample=0.060, + ) + run.setup_opt_data(phantoms=run.phantoms) + run.setup_optimiser(nwalkers=50, sample_high_density=True, model_kwargs={}) + results = run( + filepath="./results/test/", + pulse_to_write=23000101, + run="RUN01", + mds_write=True, + plot=True, + burn_frac=0.10, + iterations=100, + ) diff --git a/indica/workflows/example_bayes_opt.py b/indica/workflows/example_bayes_opt.py deleted file mode 100644 index 5e5799f4..00000000 --- a/indica/workflows/example_bayes_opt.py +++ /dev/null @@ -1,182 +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)], element="ar" - ) - 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) diff --git a/indica/workflows/load_modelling_plasma.py b/indica/workflows/load_modelling_plasma.py index e0fcd6a5..957577bc 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 @@ -31,7 +31,7 @@ DIAGNOSTIC_MODELS = { "smmh1": Interferometry, "nirh1": Interferometry, - "xrcs": Helike_spectroscopy, + "xrcs": HelikeSpectrometer, "ts": ThomsonScattering, "cxff_pi": ChargeExchange, "cxff_tws_c": ChargeExchange, 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..adf29bda 100644 --- a/indica/workflows/zeff_profile.py +++ b/indica/workflows/zeff_profile.py @@ -222,7 +222,6 @@ def prepare_inputs( phantom_data: bool = True, ts_side: str = "LFS", ): - flat_data: dict = {} models: dict = {} @@ -350,7 +349,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 +428,6 @@ def run_inversion( phantom_data: bool = True, ts_side: str = "LFS", ): - plasma, models, flat_data, input_profiles = prepare_inputs( pulse, tstart=tstart, diff --git a/indica/writers/bda_tree.py b/indica/writers/bda_tree.py new file mode 100644 index 00000000..5a5878a8 --- /dev/null +++ b/indica/writers/bda_tree.py @@ -0,0 +1,269 @@ +# Create trees from ppac standard utility tools + +import sys + +from MDSplus import Connection +import numpy as np +import standard_utility as util + + +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"), + "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"), + "TI": ("SIGNAL", "Ion temperature of main ion, eV"), + "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, "), + "ZIM3": ("SIGNAL", "Average charge of impurity IMP3, "), + "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"), + "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"), + "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"), + "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"), + }, + "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"), + "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": { + "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"), + "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"), + }, + } + 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, + mode="EDIT", +): + bda_nodes = bda() + 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 + } + + nodes = { + "RUN": ("TEXT", "RUN used for diagnostic"), + "USAGE": ("TEXT", "Quantity 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=mode, + 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 + 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) 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() 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/integration/models/test_helike_spectroscopy.py b/tests/integration/models/test_helike_spectroscopy.py index fefba5d0..770275d6 100644 --- a/tests/integration/models/test_helike_spectroscopy.py +++ b/tests/integration/models/test_helike_spectroscopy.py @@ -1,75 +1,46 @@ -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_transform_example from indica.models.plasma import example_run as example_plasma -# 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 - - class TestHelike: 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 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.HelikeSpectrometer("diagnostic_name") + self.model.set_plasma(self.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( + def test_helike_spectra_with_multiple_LOS( 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) + 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 diff --git a/tests/integration/models/test_plasma.py b/tests/integration/models/test_plasma.py new file mode 100644 index 00000000..46e9e229 --- /dev/null +++ b/tests/integration/models/test_plasma.py @@ -0,0 +1,6 @@ +import indica.models.plasma as plasma + + +def test_example_plasma(): + + _ = plasma.example_run() diff --git a/tests/integration/test_bayesmodels.py b/tests/integration/test_bayesmodels.py deleted file mode 100644 index cb11863d..00000000 --- a/tests/integration/test_bayesmodels.py +++ /dev/null @@ -1,74 +0,0 @@ -import emcee -import flatdict -from tests.integration.models.test_helike_spectroscopy import helike_LOS_example - -from indica.bayesmodels import BayesModels -from indica.bayesmodels import get_uniform -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() # pulse=9229) PHANTOM! - 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), - # "Te_prof_peaking": get_uniform(1, 5), - "Ti_prof.y0": get_uniform(2e3, 8e3), - # "Ti_prof_peaking": get_uniform(1, 5), - } - - 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", - # "Te_prof.peaking", - "Ti_prof.y0", - # "Ti_prof.peaking", - ] - - 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() 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 index 46e9e229..4746e702 100644 --- a/tests/unit/models/test_plasma.py +++ b/tests/unit/models/test_plasma.py @@ -1,6 +1,82 @@ -import indica.models.plasma as plasma +import numpy as np +from indica.equilibrium import Equilibrium +from indica.equilibrium import fake_equilibrium_data +from indica.models.plasma import Plasma -def test_example_plasma(): - _ = plasma.example_run() +class TestPlasma: + def setup_class(self): + 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.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 setup_method(self): + # self.plasma.electron_density = 1 # set profiles + # return + + def teardown_method(self): + self.plasma.initialize_variables() + + 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 + + # 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_fz_one_time_point(self): + # return + # + # def test_fz_keys_match_elements(self): + # return + + +# class TestPlasmaProfiles: +# +# def setup_class(self): +# +# def test_update_profiles +# +# def test_assign_profiles +# +# +# +# +# class TestCacheDependency: +# +# def setup_class(self): +# return +if __name__ == "__main__": + test = TestPlasma() + test.setup_class() diff --git a/tests/unit/workflows/__init__.py b/tests/unit/workflows/__init__.py new file mode 100644 index 00000000..e69de29b 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..383f19c8 --- /dev/null +++ b/tests/unit/workflows/test_bayes_workflow_example.py @@ -0,0 +1,160 @@ +import copy + +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: +Mock reader for testing experimental data reading +""" + + +class TestBayesWorkflowExample: + def setup_class(self): + self.init_settings = dict( + pulse=None, + phantoms=True, + diagnostics=["xrcs", "efit", "smmh1", "cxff_pi"], + opt_quantity=OPTIMISED_QUANTITY, + param_names=OPTIMISED_PARAMS, + profile_params=DEFAULT_PROFILE_PARAMS, + priors=DEFAULT_PRIORS, + tstart=0.02, + tend=0.10, + dt=0.005, + ) + self.plasma_settings = dict( + tsample=0.060, + ) + + self.optimiser_settings = dict( + 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, + 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 + + def setup_method(self): + self.workflow = copy.deepcopy(self.workflow_untouched) + + def teardown_method(self): + self.workflow = None + + def test_workflow_initializes(self): + 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") + assert True + + def test_init_phantoms_false_with_example_plasma(self): + with pytest.raises(ValueError): + 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})) + + # 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.plasma_settings) + assert hasattr(self.workflow.plasma, "equilibrium") + + def test_phantom_profiles_are_not_mutatable(self): + 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(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.setup_plasma(**self.plasma_settings) + self.workflow.setup_opt_data(phantoms=True) + 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.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 + + # 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.setup_plasma(**self.plasma_settings) + self.workflow.setup_opt_data(phantoms=True, noise=True) + assert self.workflow.opt_data + + 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_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.plasma_settings) + self.workflow.setup_opt_data(phantoms=True) + self.workflow.setup_optimiser(**self.optimiser_settings) + self.workflow.run_sampler(**self.sampler_settings) + if not hasattr(self.workflow, "result"): + raise ValueError("missing result in workflow object") + assert True + + +if __name__ == "__main__": + test = TestBayesWorkflowExample() + test.setup_class()