From 4dbeddd1a1cf7b80fe5800e3828944d9c7c56885 Mon Sep 17 00:00:00 2001 From: marcosertoli <74655352+marcosertoli@users.noreply.github.com> Date: Mon, 17 Jul 2023 17:58:14 +0100 Subject: [PATCH] Marcosertoli/satakihino (#263) * Temporarily stop caching (issues to be solved) * Copied Bremss bayes opt to new file --------- Co-authored-by: Marco Sertoli --- indica/models/plasma.py | 16 +- indica/workflows/example_bayes_opt.py | 89 +++++----- indica/workflows/example_bayes_opt_brems.py | 187 ++++++++++++++++++++ 3 files changed, 238 insertions(+), 54 deletions(-) create mode 100644 indica/workflows/example_bayes_opt_brems.py diff --git a/indica/models/plasma.py b/indica/models/plasma.py index 35307de7..31cfdb28 100644 --- a/indica/models/plasma.py +++ b/indica/models/plasma.py @@ -350,7 +350,6 @@ def initialize_variables(self, tstart: float, tend: float, dt: float): self._pressure_th = assign_data( self.data2d, ("pressure", "thermal"), "Pa $m^{-3}$" ) - self._ion_density = assign_data(self.data3d, ("density", "ion"), "$m^{-3}$") self._pressure_tot = assign_data( self.data2d, ("pressure", "total"), "Pa $m^{-3}$" ) @@ -598,7 +597,7 @@ def wp(self): @property def fz(self): - return self.Fz() + return self.calc_fz() # self.Fz() def calc_fz(self): for elem in self.elements: @@ -621,7 +620,7 @@ def calc_fz(self): @property def zeff(self): - return self.Zeff() + return self.calc_zeff() # Zeff() def calc_zeff(self): electron_density = self.electron_density @@ -636,7 +635,7 @@ def calc_zeff(self): @property def ion_density(self): - return self.Ion_density() + return self.calc_ion_density() # self.Ion_density() def calc_ion_density(self): meanz = self.meanz @@ -656,7 +655,7 @@ def calc_ion_density(self): @property def lz_tot(self): - return self.Lz_tot() + return self.calc_lz_tot() # self.Lz_tot() def calc_lz_tot(self): fz = self.fz @@ -681,7 +680,7 @@ def calc_lz_tot(self): @property def lz_sxr(self): - return self.Lz_sxr() + return self.calc_lz_sxr() # self.Lz_sxr() def calc_lz_sxr(self): fz = self.fz @@ -708,7 +707,7 @@ def calc_lz_sxr(self): @property def total_radiation(self): - return self.Total_radiation() + return self.calc_total_radiation() # self.Total_radiation() def calc_total_radiation(self): lz_tot = self.lz_tot @@ -728,7 +727,7 @@ def calc_total_radiation(self): @property def sxr_radiation(self): - return self.Sxr_radiation() + return self.calc_sxr_radiation() # self.Sxr_radiation() def calc_sxr_radiation(self): if not hasattr(self, "power_loss_sxr"): @@ -1238,6 +1237,7 @@ def __init__(self, operator: Callable, dependencies: list, verbose: bool = False @lru_cache() def __call__(self): + print("Recalculating") if self.verbose: print("Recalculating") return self.operator() diff --git a/indica/workflows/example_bayes_opt.py b/indica/workflows/example_bayes_opt.py index d18e04c5..4ac1bcee 100644 --- a/indica/workflows/example_bayes_opt.py +++ b/indica/workflows/example_bayes_opt.py @@ -11,8 +11,6 @@ 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 -import indica.readers.read_st40 as read_st40 -from indica.models.diode_filters import BremsstrahlungDiode # TODO: allow conditional prior usage even when only # one param is being optimisied i.e. 1 is constant @@ -34,8 +32,8 @@ def run( tend=tend, dt=dt, main_ion="h", - impurities=("c",), #impurities: tuple = ("c", "ar"), impurity_concentration: tuple = (0.02, 0.001), - impurity_concentration=(0.2,), + impurities=("ar",), + impurity_concentration=(0.001,), full_run=False, n_rad=10, ) @@ -50,27 +48,25 @@ def run( "impurity_density": plasma.Nimp_prof.yspl, } - ST40 = read_st40.ReadST40(pulse) - ST40(["pi"]) - - data_to_read=ST40.binned_data["pi"]["spectra"] - los_transform = data_to_read.transform - data_to_read.transform.set_equilibrium(data_to_read.transform.equilibrium) - pi = BremsstrahlungDiode(name="pi", channel_mask=slice(18, 28)) - pi.set_los_transform(los_transform) - - pi.plasma = plasma - pi_data = pi()["brightness"] - - #from indica.models.background_fit import Bremsstrahlung - #pi_data = Bremsstrahlung(pulse)[1] - #data_measured = Bremsstrahlung(pulse)[1].sel(channel=channels) - #data_modelled = example_run(pulse)[2]["brightness"].sel(channel=channels) + ST40 = ReadST40(pulse, tstart=tstart, tend=tend) + ST40(["xrcs", "smmh1"]) + # Initialise Diagnostic Models + los_transform = ST40.binned_data["smmh1"]["ne"].transform + smmh1 = Interferometry(name="smmh1") + smmh1.set_los_transform(los_transform) + smmh1.plasma = plasma + los_transform = ST40.binned_data["xrcs"]["te_kw"].transform + xrcs = Helike_spectroscopy(name="xrcs", window_masks=[slice(0.3945, 0.3962)]) + xrcs.set_los_transform(los_transform) + xrcs.plasma = plasma flat_data = {} - flat_data["pi.brightness"] = ( - pi_data.expand_dims(dim={"t": [plasma.time_to_calculate]}) + 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 = { @@ -93,16 +89,16 @@ def run( "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", + # "Nimp_prof.peaking", "Te_prof.y0", - # "Te_prof.peaking", + # "Te_prof.peaking", "Ti_prof.y0", # "Ti_prof.peaking", ] @@ -110,17 +106,19 @@ def run( bm = BayesModels( plasma=plasma, data=flat_data, - diagnostic_models=[pi], + diagnostic_models=[smmh1, xrcs], quant_to_optimise=[ - "pi.brightness", + "smmh1.ne", + "xrcs.spectra", ], priors=priors, ) - + ndim = param_names.__len__() - nwalkers = 50 + 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, @@ -143,22 +141,26 @@ def run( ) 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, - } + "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, @@ -175,9 +177,4 @@ def run( "Ti_prof.y0": 5000, "Ti_prof.peaking": 2, } - from sys import platform - if platform == "linux" or platform == "linux2": - pathname="./plots/" - elif platform == "win32": - pathname="C:\\Users\\Aleksandra.Alieva\\Desktop\\Plots\\New\\" - run(10607, params, 200, pathname, burn_in=0) + run(10009, params, 10, "./results/test/", burn_in=0) diff --git a/indica/workflows/example_bayes_opt_brems.py b/indica/workflows/example_bayes_opt_brems.py new file mode 100644 index 00000000..63265092 --- /dev/null +++ b/indica/workflows/example_bayes_opt_brems.py @@ -0,0 +1,187 @@ +import emcee +import numpy as np +import pandas as pd +import xarray as xr +import matplotlib.pylab as plt + +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 +import indica.readers.read_st40 as read_st40 +from indica.models.diode_filters import BremsstrahlungDiode + +# 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, + nwalkers = 10, +): + plasma = Plasma( + tstart=tstart, + tend=tend, + dt=dt, + main_ion="h", + impurities=("c",), #impurities: tuple = ("c", "ar"), impurity_concentration: tuple = (0.02, 0.001), + impurity_concentration=(0.2,), + full_run=False, + n_rad=10, + ) + plasma.build_atomic_data() + plasma.time_to_calculate = plasma.t[tsample] + + ST40 = read_st40.ReadST40(pulse) + ST40(["pi"]) + + data_to_read=ST40.binned_data["pi"]["spectra"] + los_transform = data_to_read.transform + # data_to_read.transform.set_equilibrium(data_to_read.transform.equilibrium) + pi = BremsstrahlungDiode(name="pi", channel_mask=slice(18, 28)) + pi.set_los_transform(los_transform) + pi.plasma = plasma + + # Make phantom profiles + plasma.update_profiles(phantom_profile_params) + 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, + } + + pi_data = pi()["brightness"] + + #from indica.models.background_fit import Bremsstrahlung + #pi_data = Bremsstrahlung(pulse)[1] + #data_measured = Bremsstrahlung(pulse)[1].sel(channel=channels) + #data_modelled = example_run(pulse)[2]["brightness"].sel(channel=channels) + + + flat_data = {} + flat_data["pi.brightness"] = ( + pi_data.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, 5e18), + "Nimp_prof.y1": get_uniform(1e15, 1e18), + "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), + } + + 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=[pi], + quant_to_optimise=[ + "pi.brightness", + ], + priors=priors, + ) + + ndim = param_names.__len__() + 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)) + 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__": + phantom_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": 1e18, + "Nimp_prof.y1": 1e17, + "Nimp_prof.peaking": 7, #2 + "Te_prof.y0": 3000, + "Te_prof.peaking": 2, + "Ti_prof.y0": 5000, + "Ti_prof.peaking": 2, + } + nwalkers = 50 #10 # + iterations = 200 #50 # + from sys import platform + if platform == "linux" or platform == "linux2": + pathname="./plots/" + elif platform == "win32": + pathname="C:\\Users\\Aleksandra.Alieva\\Desktop\\Plots\\New\\" + run(10607, phantom_profile_params, iterations, pathname, burn_in=0, nwalkers = nwalkers)