diff --git a/indica/bayesmodels.py b/indica/bayesmodels.py index a05e8c77..b6b5f74b 100644 --- a/indica/bayesmodels.py +++ b/indica/bayesmodels.py @@ -8,7 +8,7 @@ warnings.simplefilter("ignore", category=FutureWarning) -PROFILES = [ +PLASMA_ATTRIBUTES = [ "electron_temperature", "electron_density", "ion_temperature", @@ -17,6 +17,10 @@ "fast_density", "neutral_density", "zeff", + "wp", + "wth", + "ptot", + "pth", ] @@ -196,10 +200,11 @@ def sample_from_priors(self, param_names, size=10): def sample_from_high_density_region( self, param_names: list, sampler, nwalkers: int, nsamples=100 ): + # TODO: implement smarter MLE (maximum likelihood estimate) 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) + num_best_points = 3 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) @@ -209,7 +214,7 @@ def sample_from_high_density_region( while samples.size < param_names.__len__() * nwalkers: sample = np.random.normal( np.mean(best_start_points, axis=0), - best_points_std * 2, + best_points_std, size=(nwalkers * 5, len(param_names)), ) start = {name: sample[:, idx] for idx, name in enumerate(param_names)} @@ -249,14 +254,14 @@ def ln_posterior(self, parameters: dict, **kwargs): ln_likelihood = self._ln_likelihood() # compare results to data ln_posterior = ln_likelihood + ln_prior - plasma_profiles = {} - for profile_key in PROFILES: - if hasattr(self.plasma, profile_key): - plasma_profiles[profile_key] = getattr(self.plasma, profile_key).sel( + plasma_attributes = {} + for plasma_key in PLASMA_ATTRIBUTES: + if hasattr(self.plasma, plasma_key): + plasma_attributes[plasma_key] = getattr(self.plasma, plasma_key).sel( t=self.plasma.time_to_calculate ) else: - raise ValueError(f"plasma does not have attribute {profile_key}") + raise ValueError(f"plasma does not have attribute {plasma_key}") - blob = deepcopy({**self.bckc, **plasma_profiles}) + blob = deepcopy({**self.bckc, **plasma_attributes}) return ln_posterior, blob diff --git a/indica/models/helike_spectroscopy.py b/indica/models/helike_spectroscopy.py index 90ac5c4e..dcf50731 100644 --- a/indica/models/helike_spectroscopy.py +++ b/indica/models/helike_spectroscopy.py @@ -36,7 +36,7 @@ def __init__( name: str, instrument_method="get_helike_spectroscopy", etendue: float = 1.0, - calibration: float = 8.0e-20, + calibration: float = 1.0e-18, element: str = "ar", window_len: int = 1030, window_lim=None, @@ -204,8 +204,6 @@ def _make_spectra(self, calc_rho: bool = False): ), ) spectra = xr.concat([_spectra, empty], "wavelength") - spectra = spectra.sortby("wavelength") - self.spectra = spectra measured_spectra = self.los_transform.integrate_on_los( @@ -213,9 +211,11 @@ def _make_spectra(self, calc_rho: bool = False): t=self.spectra.t, calc_rho=calc_rho, ) - self.measured_spectra = measured_spectra.assign_coords( - {"wavelength": self.window.wavelength} + measured_spectra = measured_spectra.assign_coords( + {"wavelength": self.spectra.wavelength} ) + measured_spectra[measured_spectra==0] = np.nan + self.measured_spectra = measured_spectra.sortby("wavelength") self.spectra_los = self.los_transform.along_los def _moment_analysis(self): @@ -386,6 +386,7 @@ def __call__( calc_rho: bool = False, moment_analysis: bool = False, background: int = None, + pixel_offset: int = None, **kwargs, ): """ @@ -462,6 +463,9 @@ def __call__( if background is not None: self.measured_spectra = self.measured_spectra + background + if pixel_offset is not None: + self.measured_spectra = self.measured_spectra.shift(wavelength=round(pixel_offset), fill_value=np.nan) + self._build_bckc_dictionary() return self.bckc diff --git a/indica/models/plasma.py b/indica/models/plasma.py index 0c635f88..b119af32 100644 --- a/indica/models/plasma.py +++ b/indica/models/plasma.py @@ -419,6 +419,28 @@ def initialize_variables(self): ], ) + self.Pth = CachedCalculation( + self.calc_pth, + [ + self.electron_density, + self.ion_density, + self.electron_temperature, + self.ion_temperature, + ], + ) + + self.Ptot = CachedCalculation( + self.calc_ptot, + [ + self.electron_density, + self.ion_density, + self.electron_temperature, + self.ion_temperature, + self.pressure_fast, + ], + ) + + self.Lz_tot = CachedCalculation( self.calc_lz_tot, [ @@ -545,13 +567,7 @@ def pressure_el(self): @property def pressure_th(self): - ion_density = self.ion_density - self._pressure_th.values = self.pressure_el - for elem in self.elements: - self._pressure_th.values += ph.calc_pressure( - ion_density.sel(element=elem).values, - self.ion_temperature.sel(element=elem).values, - ) + self._pressure_th = ph.calc_pressure(self.ion_density, self.ion_temperature).sum("element") + self.pressure_el return self._pressure_th @property @@ -569,6 +585,9 @@ def pressure_fast(self): @property def pth(self): + return self.Pth() + + def calc_pth(self): pressure_th = self.pressure_th for t in np.array(self.time_to_calculate, ndmin=1): self._pth.loc[dict(t=t)] = np.trapz( @@ -578,6 +597,9 @@ def pth(self): @property def ptot(self): + return self.Ptot() + + def calc_ptot(self): pressure_tot = self.pressure_tot for t in np.array(self.time_to_calculate, ndmin=1): self._ptot.loc[dict(t=t)] = np.trapz( @@ -599,7 +621,8 @@ def wp(self): @property def fz(self): - return self.calc_fz() # self.Fz() + return self.Fz() + # return self.calc_fz() # self.Fz() def calc_fz(self): for elem in self.elements: @@ -622,37 +645,29 @@ def calc_fz(self): @property def zeff(self): - return self.calc_zeff() # Zeff() + return self.Zeff() + # return self.calc_zeff() def calc_zeff(self): - electron_density = self.electron_density - ion_density = self.ion_density - meanz = self.meanz - for elem in self.elements: - self._zeff.loc[dict(element=elem)] = ( - (ion_density.sel(element=elem) * meanz.sel(element=elem) ** 2) - / electron_density - ).values + self._zeff = (self.ion_density * self.meanz) ** 2 / self.electron_density return self._zeff @property def ion_density(self): - return self.calc_ion_density() # self.Ion_density() + return self.Ion_density() + # return self.calc_ion_density() def calc_ion_density(self): meanz = self.meanz - main_ion_density = self.electron_density - self.fast_density * meanz.sel( - element=self.main_ion - ) for elem in self.impurities: self._ion_density.loc[dict(element=elem)] = self.impurity_density.sel( element=elem ).values - main_ion_density -= self.impurity_density.sel(element=elem) * meanz.sel( - element=elem - ) - self._ion_density.loc[dict(element=self.main_ion)] = main_ion_density.values + main_ion_density = self.electron_density - self.fast_density * meanz.sel(element=self.main_ion) \ + - (self.impurity_density * meanz).sum("element") + + self._ion_density.loc[dict(element=self.main_ion)] = main_ion_density return self._ion_density @property @@ -753,6 +768,7 @@ def calc_sxr_radiation(self): @property def meanz(self): fz = self.fz + for elem in self.elements: self._meanz.loc[dict(element=elem)] = ( (fz[elem] * fz[elem].ion_charges).sum("ion_charges").values @@ -1234,7 +1250,6 @@ 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/models/thomson_scattering.py b/indica/models/thomson_scattering.py index d74498ee..94af48e8 100644 --- a/indica/models/thomson_scattering.py +++ b/indica/models/thomson_scattering.py @@ -41,6 +41,9 @@ def _build_bckc_dictionary(self): self.bckc[quantity] = self.Te_at_channels long_name = "Te" units = "eV" + elif quant == "chi2": + # Placeholder + continue else: print(f"{quant} not available in model for {self.instrument_method}") continue @@ -82,7 +85,7 @@ def __call__( """ if self.plasma is not None: if t is None: - t = self.plasma.t + t = self.plasma.time_to_calculate Ne = self.plasma.electron_density.interp(t=t) Te = self.plasma.electron_temperature.interp(t=t) else: @@ -111,6 +114,18 @@ def __call__( return self.bckc +def ts_transform_example(nchannels): + 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) + transform = TransectCoordinates( + x_positions, + y_positions, + z_positions, + "ts", + machine_dimensions=((0.15, 0.95), (-0.7, 0.7)), + ) + return transform def example_run( pulse: int = None, @@ -124,19 +139,7 @@ def example_run( if plasma is None: plasma = example_plasma(pulse=pulse) - # Create new interferometers diagnostics - nchannels = 11 - 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 = ts_transform_example(11) transect_transform.set_equilibrium(plasma.equilibrium) model = ThomsonScattering( diagnostic_name, diff --git a/indica/profiles_gauss.py b/indica/profiles_gauss.py index e266a4ed..67b1cf7e 100644 --- a/indica/profiles_gauss.py +++ b/indica/profiles_gauss.py @@ -32,6 +32,7 @@ def __init__( self.xend = xend self.coord = f"rho_{coord}" self.x = np.linspace(0, 1, 15) ** 0.7 + # self.x = 1-np.logspace(0, -2, 30) self.datatype = datatype if xspl is None: xspl = np.linspace(0, 1.0, 30) @@ -55,6 +56,9 @@ def __init__( if parameters is None: parameters = get_defaults(datatype) + elif {"y0", "y1", "yend", "wcenter", "wped", "peaking",} >= set(parameters): + _parameters = get_defaults(datatype) + parameters = dict(_parameters, **parameters) for k, p in parameters.items(): setattr(self, k, p) diff --git a/indica/readers/read_st40.py b/indica/readers/read_st40.py index 62cdaf62..f14f999e 100644 --- a/indica/readers/read_st40.py +++ b/indica/readers/read_st40.py @@ -120,8 +120,8 @@ def get_raw_data(self, uid: str, instrument: str, revision: RevisionLike = 0): continue transform = data[quant].transform - if hasattr(transform, "set_equilibrium"): - transform.set_equilibrium(self.equilibrium) + # if hasattr(transform, "set_equilibrium"): + # transform.set_equilibrium(self.equilibrium) self.transforms[instrument] = transform self.raw_data[instrument] = data diff --git a/indica/workflows/abstract_bayes_workflow.py b/indica/workflows/abstract_bayes_workflow.py index 4839a2b5..24882966 100644 --- a/indica/workflows/abstract_bayes_workflow.py +++ b/indica/workflows/abstract_bayes_workflow.py @@ -47,6 +47,11 @@ def read_test_data( 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) + + missing_keys = set(diagnostics) - set(self.reader.binned_data.keys()) + if len(missing_keys) > 0: + raise ValueError(f"missing data: {missing_keys}") + self.equilibrium = self.reader.equilibrium self.transforms = self.reader.transforms self.data = self.reader.binned_data @@ -127,21 +132,21 @@ def save_phantom_profiles(self, kinetic_profiles=None): self.phantom_profiles = phantom_profiles - def _build_result_dict(self): + def _build_inputs_dict(self): """ Returns ------- - dictionary of results in MDS+ structure + dictionary of inputs 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["TIME_BINS"] = self.plasma.t + result["TIME"] = self.plasma.time_to_calculate result["METADATA"] = { "GITCOMMIT": "PLACEHOLDER", @@ -157,9 +162,13 @@ def _build_result_dict(self): "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, + "TSTART":self.tstart, + "TEND": self.tend, + "DT": self.dt, + "TSAMPLE": self.tsample, + } result["INPUT"]["WORKFLOW"] = { diag_name.upper(): { @@ -172,24 +181,33 @@ def _build_result_dict(self): for diag_name in self.diagnostics } - result["MODEL_DATA"] = { + result["DIAG_DATA"] = { diag_name.upper(): { - quantity[1].upper(): self.blobs[f"{quantity[0]}.{quantity[1]}"] + 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["MODEL_DATA"]["SAMPLES"] = self.samples - result["DIAG_DATA"] = { + + return result + + + def _build_result_dict(self, ): + result = {} + quant_list = [item.split(".") for item in self.opt_quantity] + + result["MODEL_DATA"] = { diag_name.upper(): { - quantity[1].upper(): self.opt_data[f"{quantity[0]}.{quantity[1]}"] + 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["PHANTOMS"] = { "FLAG": self.phantoms, @@ -226,31 +244,33 @@ def _build_result_dict(self): "RHO_TOR": self.plasma.equilibrium.rhotor.interp(t=self.plasma.t), "NE": self.blobs["electron_density"].median(dim="index"), "NI": self.blobs["ion_density"] - .sel(element=self.plasma.main_ion) - .median(dim="index"), + .sel(element=self.plasma.main_ion) + .median(dim="index"), "TE": self.blobs["electron_temperature"].median(dim="index"), "TI": self.blobs["ion_temperature"] - .sel(element=self.plasma.main_ion) - .median(dim="index"), + .sel(element=self.plasma.main_ion) + .median(dim="index"), "NFAST": self.blobs["fast_density"].median(dim="index"), "NNEUTR": self.blobs["neutral_density"].median(dim="index"), "NE_ERR": self.blobs["electron_density"].std(dim="index"), "NI_ERR": self.blobs["ion_density"] - .sel(element=self.plasma.main_ion) - .std(dim="index"), + .sel(element=self.plasma.main_ion) + .std(dim="index"), "TE_ERR": self.blobs["electron_temperature"].std(dim="index"), "TI_ERR": self.blobs["ion_temperature"] - .sel(element=self.plasma.main_ion) - .std(dim="index"), + .sel(element=self.plasma.main_ion) + .std(dim="index"), "NFAST_ERR": self.blobs["fast_density"].std(dim="index"), "NNEUTR_ERR": self.blobs["neutral_density"].std(dim="index"), + "ZEFF": self.blobs["zeff"].sum("element").median(dim="index"), + "ZEFF_ERR": self.blobs["zeff"].sum("element").std(dim="index"), } result["PROFILES"] = { **result["PROFILES"], **{ f"NIZ{num_imp + 1}": self.blobs["impurity_density"] - .sel(element=imp) - .median(dim="index") + .sel(element=imp) + .median(dim="index") for num_imp, imp in enumerate(self.plasma.impurities) }, } @@ -258,8 +278,8 @@ def _build_result_dict(self): **result["PROFILES"], **{ f"NIZ{num_imp + 1}_ERR": self.blobs["impurity_density"] - .sel(element=imp) - .std(dim="index") + .sel(element=imp) + .std(dim="index") for num_imp, imp in enumerate(self.plasma.impurities) }, } @@ -267,8 +287,8 @@ def _build_result_dict(self): **result["PROFILES"], **{ f"TIZ{num_imp + 1}": self.blobs["ion_temperature"] - .sel(element=imp) - .median(dim="index") + .sel(element=imp) + .median(dim="index") for num_imp, imp in enumerate(self.plasma.impurities) }, } @@ -276,8 +296,8 @@ def _build_result_dict(self): **result["PROFILES"], **{ f"TIZ{num_imp + 1}_ERR": self.blobs["ion_temperature"] - .sel(element=imp) - .std(dim="index") + .sel(element=imp) + .std(dim="index") for num_imp, imp in enumerate(self.plasma.impurities) }, } @@ -312,45 +332,48 @@ def _build_result_dict(self): "PRIOR_SAMPLE": self.prior_sample, "POST_SAMPLE": self.post_sample, "AUTOCORR": self.autocorr, + "GELMANRUBIN": gelman_rubin(self.sampler.get_chain(flat=False)) } result["GLOBAL"] = { "TI0": self.blobs["ion_temperature"] - .sel(element=self.plasma.main_ion) - .sel(rho_poloidal=0, method="nearest") - .median(dim="index"), + .sel(element=self.plasma.main_ion) + .sel(rho_poloidal=0, method="nearest") + .median(dim="index"), "TE0": self.blobs["electron_temperature"] - .sel(rho_poloidal=0, method="nearest") - .median(dim="index"), + .sel(rho_poloidal=0, method="nearest") + .median(dim="index"), "NE0": self.blobs["electron_density"] - .sel(rho_poloidal=0, method="nearest") - .median(dim="index"), + .sel(rho_poloidal=0, method="nearest") + .median(dim="index"), "NI0": self.blobs["ion_density"] - .sel(element=self.plasma.main_ion) - .sel(rho_poloidal=0, method="nearest") - .median(dim="index"), - "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"), + .sel(element=self.plasma.main_ion) + .sel(rho_poloidal=0, method="nearest") + .median(dim="index"), + "WP": self.blobs["wp"] + .median(dim="index"), + "WP_ERR": self.blobs["wp"] + .std(dim="index"), + "WTH": self.blobs["wth"] + .median(dim="index"), + "WTH_ERR": self.blobs["wth"] + .std(dim="index"), + "PTOT": self.blobs["ptot"] + .median(dim="index"), + "PTOT_ERR": self.blobs["ptot"] + .std(dim="index"), + "PTH": self.blobs["pth"] + .median(dim="index"), + "PTH_ERR": self.blobs["pth"] + .std(dim="index"), } result["GLOBAL"] = { **result["GLOBAL"], **{ f"TI0Z{num_imp + 1}": self.blobs["ion_temperature"] - .sel(element=imp) - .sel(rho_poloidal=0, method="nearest") - .median(dim="index") + .sel(element=imp) + .sel(rho_poloidal=0, method="nearest") + .median(dim="index") for num_imp, imp in enumerate(self.plasma.impurities) }, } @@ -358,9 +381,9 @@ def _build_result_dict(self): **result["GLOBAL"], **{ f"TI0Z{num_imp + 1}_ERR": self.blobs["ion_temperature"] - .sel(element=imp) - .sel(rho_poloidal=0, method="nearest") - .std(dim="index") + .sel(element=imp) + .sel(rho_poloidal=0, method="nearest") + .std(dim="index") for num_imp, imp in enumerate(self.plasma.impurities) }, } @@ -368,9 +391,9 @@ def _build_result_dict(self): **result["GLOBAL"], **{ f"NI0Z{num_imp + 1}": self.blobs["impurity_density"] - .sel(element=imp) - .sel(rho_poloidal=0, method="nearest") - .median(dim="index") + .sel(element=imp) + .sel(rho_poloidal=0, method="nearest") + .median(dim="index") for num_imp, imp in enumerate(self.plasma.impurities) }, } @@ -378,15 +401,13 @@ def _build_result_dict(self): **result["GLOBAL"], **{ f"NI0Z{num_imp + 1}_ERR": self.blobs["impurity_density"] - .sel(element=imp) - .sel(rho_poloidal=0, method="nearest") - .std(dim="index") + .sel(element=imp) + .sel(rho_poloidal=0, method="nearest") + .std(dim="index") for num_imp, imp in enumerate(self.plasma.impurities) }, } - - self.result = result - return self.result + return result def run_sampler(self, iterations, burn_frac): """ @@ -409,10 +430,13 @@ def run_sampler(self, iterations, burn_frac): self.param_names.__len__(), auto_sample=10, ) - blobs = self.sampler.get_blobs(discard=int(iterations * burn_frac), flat=True) - blob_names = self.sampler.get_blobs().flatten()[0].keys() + _blobs = self.sampler.get_blobs(discard=int(iterations * burn_frac), flat=True) + blobs = [blob for blob in _blobs if blob] # remove empty blobs + + blob_names = blobs[0].keys() self.samples = np.arange(0, blobs.__len__()) + self.blobs = { blob_name: xr.concat( [data[blob_name] for data in blobs], @@ -424,14 +448,13 @@ def run_sampler(self, iterations, burn_frac): 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() + self.post_sample = self.sampler.get_chain(discard=int(iterations * burn_frac), flat=True) - def save_pickle(self, filepath): + def save_pickle(self, result, filepath): if filepath: Path(filepath).mkdir(parents=True, exist_ok=True) with open(filepath + "results.pkl", "wb") as handle: - pickle.dump(self.result, handle) + pickle.dump(result, handle) def dict_of_dataarray_to_numpy(self, dict_of_dataarray): """ @@ -454,13 +477,14 @@ def __call__(self, filepath="./results/test/", **kwargs): return self.result -def sample_with_autocorr(sampler, start_points, iterations, n_params, auto_sample=5): +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, + skip_initial_state_check=True, ): if sampler.iteration % auto_sample: continue @@ -477,3 +501,15 @@ def sample_with_autocorr(sampler, start_points, iterations, n_params, auto_sampl : sampler.iteration, ] return autocorr + +def gelman_rubin(chain): + ssq = np.var(chain, axis=1, ddof=1) + w = np.mean(ssq, axis=0) + theta_b = np.mean(chain, axis=1) + theta_bb = np.mean(theta_b, axis=0) + m = chain.shape[0] + n = chain.shape[1] + B = n/(m-1) * np.sum((theta_bb - theta_b)**2, axis=0) + var_theta = (n-1) / n * w + 1 / n * B + R = np.sqrt(var_theta / w) + return R diff --git a/indica/workflows/batch_bda.py b/indica/workflows/batch_bda.py new file mode 100644 index 00000000..ef698ec1 --- /dev/null +++ b/indica/workflows/batch_bda.py @@ -0,0 +1,287 @@ +from indica.workflows.bayes_workflow_example import BayesWorkflowExample, DEFAULT_PRIORS, DEFAULT_PROFILE_PARAMS + +PARAMS_SET_TS =[ + # "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", + "xrcs.pixel_offset", + ] + +PARAMS_DEFAULT =[ + "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", + "xrcs.pixel_offset", + ] + +PARAMS_SET_ALL =[ + # "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", + "xrcs.pixel_offset", + ] + + +DIAG_DEFAULT = [ + "xrcs", + "ts", + "efit", + "cxff_pi", + "cxff_tws_c" + # "smmh1", + ] + +DIAG_DEFAULT_CHERS = [ + "xrcs", + "ts", + "efit", + # "cxff_pi", + "cxff_tws_c" + # "smmh1", + ] + +DIAG_DEFAULT_PI = [ + "xrcs", + "ts", + "efit", + "cxff_pi", + # "cxff_tws_c" + # "smmh1", + ] + + +DIAG_SET_TS = [ + "xrcs", + # "ts", + "efit", + "cxff_pi", + "cxff_tws_c" + # "smmh1", + ] + +OPT_DEFAULT = [ + "xrcs.spectra", + "ts.ne", + "ts.te", + "efit.wp", + "cxff_pi.ti", + "cxff_tws_c.ti", + # "smmh1.ne" + ] + +OPT_SET_TS = [ + "xrcs.spectra", + # "ts.ne", + # "ts.te", + "efit.wp", + "cxff_pi.ti", + "cxff_tws_c.ti", + # "smmh1.ne" + ] + +def pulse_example(pulse, diagnostics, param_names, opt_quantity, t_opt, iterations=100, run="RUN01", + fast_particles=False, astra_run="RUN602", astra_pulse_range=13000000, + astra_equilibrium=False, efit_revision=0, set_ts_profiles=False, set_all_profiles=True, + astra_wp=False, **kwargs): + + workflow = BayesWorkflowExample( + pulse=pulse, + diagnostics=diagnostics, + param_names=param_names, + opt_quantity=opt_quantity, + priors=DEFAULT_PRIORS, + profile_params=DEFAULT_PROFILE_PARAMS, + phantoms=False, + fast_particles=fast_particles, + tstart=0.06, + tend=0.10, + dt=0.005, + astra_run=astra_run, + astra_pulse_range=astra_pulse_range, + astra_equilibrium=astra_equilibrium, + efit_revision=efit_revision, + set_ts_profiles = set_ts_profiles, + set_all_profiles=set_all_profiles, + astra_wp = astra_wp, + ) + + workflow.setup_plasma( + tsample=t_opt, + # n_rad=50 + ) + workflow.setup_opt_data(phantoms=workflow.phantoms) + workflow.setup_optimiser(nwalkers=40, sample_method="high_density", model_kwargs=workflow.model_kwargs, nsamples=100) + results = workflow( + filepath=f"./results/{workflow.pulse}.{run}/", + pulse_to_write=25000000 + workflow.pulse, + run=run, + mds_write=False, + plot=True, + burn_frac=0.20, + iterations=iterations, + ) + +# (11089, DIAG_DEFAULT_CHERS, PARAMS_DEFAULT, [ + # "xrcs.spectra", + # "ts.ne", + # "ts.te", + # "efit.wp", + # # "cxff_pi.ti", + # "cxff_tws_c.ti", + # # "smmh1.ne" + # ], 0.100, + # dict(run="RUN01", fast_particles=True, astra_run="RUN601", astra_pulse_range=13000000, astra_equilibrium=False, + # efit_revision=2)), + + + +if __name__ == "__main__": + + pulse_info = [ + + # (11211, DIAG_DEFAULT, PARAMS_DEFAULT, OPT_DEFAULT, 0.084), + # (11215, DIAG_DEFAULT_CHERS, PARAMS_DEFAULT, [ + # "xrcs.spectra", + # "ts.ne", + # "ts.te", + # "efit.wp", + # # "cxff_pi.ti", + # "cxff_tws_c.ti", + # # "smmh1.ne" + # ], 0.070, dict(fast_particles=True, astra_run="RUN603", astra_pulse_range=13000000, astra_equilibrium=True, astra_wp=True)), + # + # # (11224, DIAG_DEFAULT, PARAMS_DEFAULT, OPT_DEFAULT, 0.090), + # (11225, DIAG_DEFAULT_PI, PARAMS_DEFAULT, [ + # "xrcs.spectra", + # "ts.ne", + # "ts.te", + # "efit.wp", + # "cxff_pi.ti", + # # "cxff_tws_c.ti", + # # "smmh1.ne" + # ], 0.090, dict(efit_revision=2, fast_particles=True, astra_run="RUN603", astra_pulse_range=13000000, astra_equilibrium=True, astra_wp=True)), + # + # # (11226, DIAG_DEFAULT, PARAMS_DEFAULT, OPT_DEFAULT, 0.070), + # (11227, DIAG_DEFAULT_PI, PARAMS_DEFAULT, [ + # "xrcs.spectra", + # "ts.ne", + # "ts.te", + # "efit.wp", + # "cxff_pi.ti", + # # "cxff_tws_c.ti", + # # "smmh1.ne" + # ], 0.070, dict(fast_particles=True, astra_run="RUN603", astra_pulse_range=13000000, astra_equilibrium=True, astra_wp=True)), + # + # (11228, DIAG_DEFAULT_PI, PARAMS_DEFAULT, [ + # "xrcs.spectra", + # "ts.ne", + # "ts.te", + # "efit.wp", + # "cxff_pi.ti", + # # "cxff_tws_c.ti", + # # "smmh1.ne" + # ], 0.080, dict(run="RUN02", fast_particles=True, astra_run="RUN603", astra_pulse_range=13000000, astra_equilibrium=True, astra_wp=True)), + # (11238, DIAG_DEFAULT_PI, PARAMS_DEFAULT, [ + # "xrcs.spectra", + # "ts.ne", + # "ts.te", + # "efit.wp", + # "cxff_pi.ti", + # # "cxff_tws_c.ti", + # # "smmh1.ne" + # ], 0.075, dict(efit_revision=2, fast_particles=True, astra_run="RUN603", astra_pulse_range=13000000, astra_equilibrium=True, astra_wp=True)), + # + # (11312, DIAG_DEFAULT_PI, PARAMS_DEFAULT, [ + # "xrcs.spectra", + # "ts.ne", + # "ts.te", + # "efit.wp", + # "cxff_pi.ti", + # # "cxff_tws_c.ti", + # # "smmh1.ne" + # ], 0.080, dict(run="RUN01", fast_particles=True, astra_run="RUN14", astra_pulse_range=33000000, astra_equilibrium=True, set_ts_profiles=False, set_all_profiles=False, astra_wp=True)), + # + # (11314, DIAG_DEFAULT_PI, PARAMS_DEFAULT, [ + # "xrcs.spectra", + # "ts.ne", + # "ts.te", + # "efit.wp", + # "cxff_pi.ti", + # # "cxff_tws_c.ti", + # # "smmh1.ne" + # ], 0.080, dict(run="RUN01", fast_particles=True, astra_run="RUN3", astra_pulse_range=33000000, astra_equilibrium=True, set_ts_profiles=False, astra_wp=True)), + (11312, DIAG_SET_TS, PARAMS_SET_TS, [ + "xrcs.spectra", + # "ts.ne", + # "ts.te", + # "efit.wp", + "cxff_pi.ti", + # "cxff_tws_c.ti", + # "smmh1.ne" + ], 0.080, + dict(run="TEST_TI", fast_particles=True, astra_run="RUN14", astra_pulse_range=33000000, astra_equilibrium=True, + set_ts_profiles=True, set_all_profiles=False, astra_wp=True)), + + # Tree doesnt exist + ## (11317, 0.080, dict(run="RUN01", fast_particles=True, astra_run="RUN11", astra_pulse_range=33000000, astra_equilibrium=True, set_ts_profiles=True, astra_wp=True)), + ] + + # for info in pulse_info: + # if len(info) < 6: + # info = list(info) + # info.append(dict()) + # info = tuple(info) + # try: + # pulse_example(*info[0:5], iterations=2, **info[5]) + # except Exception as e: + # print(f"pulse: {info[0]}") + # print(repr(e)) + # continue + + for info in pulse_info: + pulse_example(*info[0:5], iterations=2000, **info[5]) \ No newline at end of file diff --git a/indica/workflows/bayes_plots.py b/indica/workflows/bayes_plots.py index 93955202..15d61da1 100644 --- a/indica/workflows/bayes_plots.py +++ b/indica/workflows/bayes_plots.py @@ -45,8 +45,8 @@ def plot_profile( ) plt.fill_between( profile.rho_poloidal, - profile.quantile(0.00, dim="index"), - profile.quantile(1.00, dim="index"), + profile.quantile(0.005, dim="index"), + profile.quantile(0.995, dim="index"), label=f"{blobkey}, Max-Min", zorder=1, color="lightgrey", @@ -92,6 +92,7 @@ def _plot_1d( xlabel="[]", xlim=None, figsize=(6.4, 4.8), + hide_legend=False, **kwargs, ): if blobkey not in blobs.keys(): @@ -118,8 +119,8 @@ def _plot_1d( ) plt.fill_between( blob_data.__getattr__(dims[0]), - blob_data.quantile(0.00, dim="index"), - blob_data.quantile(1.00, dim="index"), + blob_data.quantile(0.005, dim="index"), + blob_data.quantile(0.995, dim="index"), label=f"{blobkey}, Max-Min", zorder=1, color="lightgrey", @@ -140,11 +141,14 @@ def _plot_1d( label=f"{blobkey} data", zorder=4, ) + + plt.gca().set_ylim(bottom=0) set_axis_sci() plt.ylabel(ylabel) plt.xlabel(xlabel) plt.xlim(xlim) - plt.legend() + if not hide_legend: + plt.legend() plt.savefig(figheader + filename) plt.close() @@ -181,8 +185,8 @@ def violinplot( 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) + # bot = axs.get_ylim()[0] + axs.set_ylim(top=top * 1.1, bottom=0) axs.set_ylabel(f"{ylabel}") set_axis_sci() @@ -263,10 +267,14 @@ def plot_bayes_result( plot_autocorr(autocorr, param_names, figheader, filetype=filetype) if "CXFF_PI.TI" in model_data.keys(): + max_channel = diag_data["CXFF_PI.TI"].sel(t=model_data["CXFF_PI.TI"].t).idxmax("channel").values model_data["CXFF_PI.TI0"] = model_data["CXFF_PI.TI"].sel( - channel=diag_data["CXFF_PI.TI"].channel - ) + channel=max_channel + ) diag_data["CXFF_PI.TI0"] = diag_data["CXFF_PI.TI"].sel( + channel=max_channel + ) + model_data["CXFF_PI.TI0"] = model_data["CXFF_PI.TI"].sel( channel=diag_data["CXFF_PI.TI"].channel ) @@ -346,6 +354,43 @@ def plot_bayes_result( figheader=figheader, ylabel="Temperature [eV]", xlabel="Channel", + hide_legend=True, + ) + key = "CXFF_TWS_C.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", + hide_legend=True, + + ) + + key = "TS.TE" + 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 = "TS.NE" + if key in model_data.keys(): + _plot_1d( + model_data, + key, + diag_data, + f"{key.replace('.', '_')}" + filetype, + figheader=figheader, + ylabel="Density [m^-3]", + xlabel="Channel", ) key = "TE" diff --git a/indica/workflows/bayes_workflow_example.py b/indica/workflows/bayes_workflow_example.py index cb81a09d..fa6b06ea 100644 --- a/indica/workflows/bayes_workflow_example.py +++ b/indica/workflows/bayes_workflow_example.py @@ -1,10 +1,11 @@ -from typing import Any +from typing import Any, List, Tuple from typing import Dict import emcee import flatdict import numpy as np from scipy.stats import loguniform +import xarray as xr from indica.bayesmodels import BayesModels from indica.bayesmodels import get_uniform @@ -15,84 +16,98 @@ from indica.models.helike_spectroscopy import HelikeSpectrometer from indica.models.interferometry import Interferometry from indica.models.interferometry import smmh1_transform_example +from indica.models.thomson_scattering import ThomsonScattering +from indica.models.thomson_scattering import ts_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 +from indica.readers.read_st40 import ReadST40 +from indica.equilibrium import Equilibrium + + # 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, + "Ne_prof.wped": 3, + "Ne_prof.wcenter": 0.3, + "Ne_prof.peaking": 1.2, + "Nimp_prof.y0": 1e17, - "Nimp_prof.y1": 5e15, - "Nimp_prof.wcenter": 0.4, - "Nimp_prof.wped": 6, + "Nimp_prof.y1": 1e15, + "Nimp_prof.yend": 1e15, + "Nimp_prof.wcenter": 0.3, + "Nimp_prof.wped": 3, "Nimp_prof.peaking": 2, + "Te_prof.y0": 3000, "Te_prof.y1": 50, - "Te_prof.wcenter": 0.4, + "Te_prof.yend": 10, + "Te_prof.wcenter": 0.2, "Te_prof.wped": 3, - "Te_prof.peaking": 2, + "Te_prof.peaking": 1.5, + "Ti_prof.y0": 6000, "Ti_prof.y1": 50, - "Ti_prof.wcenter": 0.4, + "Ti_prof.yend": 10, + "Ti_prof.wcenter": 0.2, "Ti_prof.wped": 3, - "Ti_prof.peaking": 2, + "Ti_prof.peaking": 1.5, } DEFAULT_PRIORS = { "Ne_prof.y0": get_uniform(2e19, 4e20), - "Ne_prof.y1": get_uniform(1e18, 1e19), + "Ne_prof.y1": get_uniform(1e18, 2e19), "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.wped": get_uniform(2, 6), + "Ne_prof.wcenter": get_uniform(0.2, 0.4), + "Ne_prof.peaking": get_uniform(1, 4), + "Nimp_prof.y0": loguniform(1e15, 1e18), + "Nimp_prof.y1": loguniform(1e14, 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.wped": get_uniform(2, 6), + "Nimp_prof.wcenter": get_uniform(0.2, 0.4), + "Nimp_prof.peaking": get_uniform(1, 6), "Nimp_prof.peaking/Ne_prof.peaking": lambda x1, x2: np.where( (x1 > x2), 1, 0 ), # impurity always more peaked + "Te_prof.y0": get_uniform(1000, 5000), - "Te_prof.wped": get_uniform(1, 6), - "Te_prof.wcenter": get_uniform(0.1, 0.8), - "Te_prof.peaking": get_uniform(1, 10), + "Te_prof.wped": get_uniform(2, 6), + "Te_prof.wcenter": get_uniform(0.2, 0.4), + "Te_prof.peaking": get_uniform(1, 4), # "Ti_prof.y0/Te_prof.y0": lambda x1, x2: np.where(x1 > x2, 1, 0), # hot ion mode "Ti_prof.y0": get_uniform(1000, 10000), - "Ti_prof.wped": get_uniform(1, 6), - "Ti_prof.wcenter": get_uniform(0.1, 0.8), - "Ti_prof.peaking": get_uniform(1, 10), + "Ti_prof.wped": get_uniform(2, 6), + "Ti_prof.wcenter": get_uniform(0.2, 0.4), + "Ti_prof.peaking": get_uniform(1, 6), + "xrcs.pixel_offset": get_uniform(-4.01, -4.0), } OPTIMISED_PARAMS = [ - # "Ne_prof.y1", + "Ne_prof.y1", "Ne_prof.y0", - # "Ne_prof.peaking", + "Ne_prof.peaking", # "Ne_prof.wcenter", - # "Ne_prof.wped", + "Ne_prof.wped", # "Nimp_prof.y1", "Nimp_prof.y0", # "Nimp_prof.wcenter", # "Nimp_prof.wped", - # "Nimp_prof.peaking", + "Nimp_prof.peaking", "Te_prof.y0", - # "Te_prof.wped", + "Te_prof.wped", "Te_prof.wcenter", "Te_prof.peaking", "Ti_prof.y0", - # "Ti_prof.wped", + "Ti_prof.wped", "Ti_prof.wcenter", "Ti_prof.peaking", ] @@ -100,7 +115,9 @@ "xrcs.spectra", "cxff_pi.ti", "efit.wp", - "smmh1.ne", + # "smmh1.ne", + "ts.te", + "ts.ne", ] @@ -113,10 +130,19 @@ def __init__( priors: dict, profile_params: dict, pulse: int = None, - phantoms: bool = False, tstart=0.02, tend=0.10, dt=0.005, + + phantoms: bool = False, + fast_particles = False, + astra_run=None, + astra_pulse_range=13000000, + astra_equilibrium=False, + efit_revision = 0, + set_ts_profiles = False, + set_all_profiles=False, + astra_wp = False, ): self.pulse = pulse self.diagnostics = diagnostics @@ -124,11 +150,22 @@ def __init__( 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 + self.phantoms = phantoms + self.fast_particles = fast_particles + self.astra_run= astra_run + self.astra_pulse_range = astra_pulse_range + self.astra_equilibrium = astra_equilibrium + self.efit_revision = efit_revision + self.set_ts_profiles = set_ts_profiles + self.set_all_profiles = set_all_profiles + self.astra_wp = astra_wp + + self.model_kwargs = {} + for attribute in [ "param_names", "opt_quantity", @@ -151,6 +188,7 @@ def __init__( "xrcs": helike_transform_example(1), "smmh1": smmh1_transform_example(1), "cxff_pi": pi_transform_example(5), + "ts":ts_transform_example(11), } self.read_test_data( example_transforms, tstart=self.tstart, tend=self.tend, dt=self.dt @@ -159,6 +197,10 @@ def __init__( self.read_data( self.diagnostics, tstart=self.tstart, tend=self.tend, dt=self.dt ) + if self.efit_revision != 0: + efit_equilibrium = self.reader.reader_equil.get("", "efit", self.efit_revision) + self.equilibrium = Equilibrium(efit_equilibrium) + self.setup_models(self.diagnostics) def setup_plasma( @@ -166,7 +208,7 @@ def setup_plasma( tstart=None, tend=None, dt=None, - tsample=0.050, + tsample=None, main_ion="h", impurities=("ar", "c"), impurity_concentration=(0.001, 0.04), @@ -189,12 +231,21 @@ def setup_plasma( 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() - ] + + if tsample == None: + self.tsample = self.plasma.t + else: + self.tsample = self.plasma.t[ + np.abs(tsample - self.plasma.t).argmin() + ] + + self.plasma.time_to_calculate = self.tsample self.plasma.set_equilibrium(self.equilibrium) self.plasma.update_profiles(self.profile_params) + if self.fast_particles: + self._init_fast_particles(run=self.astra_run) + self.plasma.update_profiles({}) + self.plasma.build_atomic_data(calc_power_loss=False) self.save_phantom_profiles() @@ -219,13 +270,13 @@ def setup_models(self, diagnostics: list): # machine_dimensions=machine_dims, # passes=2, # ) - los_transform.set_equilibrium(self.equilibrium) + los_transform.set_equilibrium(self.equilibrium, force=True) model = Interferometry(name=diag) model.set_los_transform(los_transform) elif diag == "xrcs": los_transform = self.transforms[diag] - los_transform.set_equilibrium(self.equilibrium) + los_transform.set_equilibrium(self.equilibrium, force=True) window = None if hasattr(self, "data"): if diag in self.data.keys(): @@ -243,13 +294,74 @@ def setup_models(self, diagnostics: list): elif diag == "cxff_pi": transform = self.transforms[diag] - transform.set_equilibrium(self.equilibrium) + transform.set_equilibrium(self.equilibrium, force=True) model = ChargeExchange(name=diag, element="ar") model.set_transect_transform(transform) + + elif diag == "cxff_tws_c": + transform = self.transforms[diag] + transform.set_equilibrium(self.equilibrium, force=True) + model = ChargeExchange(name=diag, element="c") + model.set_transect_transform(transform) + + elif diag == "ts": + transform = self.transforms[diag] + transform.set_equilibrium(self.equilibrium, force=True) + model = ThomsonScattering(name=diag, ) + model.set_transect_transform(transform) else: raise ValueError(f"{diag} not found in setup_models") self.models[diag] = model + def _init_fast_particles(self, run="RUN602", ): + + st40_code = ReadST40(self.astra_pulse_range + self.pulse, self.tstart-self.dt, self.tend+self.dt, dt=self.dt, tree="astra") + astra_data = st40_code.get_raw_data("", "astra", run) + + if self.astra_equilibrium: + self.equilibrium = Equilibrium(astra_data) + self.plasma.equilibrium = self.equilibrium + + st40_code.bin_data_in_time(["astra"], self.tstart, self.tend, self.dt) + code_data = st40_code.binned_data["astra"] + Nf = ( + code_data["nf"].interp(rho_poloidal=self.plasma.rho, t=self.plasma.t) + * 1.0e19 + ) + self.plasma.fast_density.values = Nf.values + Nn = ( + code_data["nn"].interp(rho_poloidal=self.plasma.rho, t=self.plasma.t) + * 1.0e19 + ) + self.plasma.neutral_density.values = Nn.values + Pblon = code_data["pblon"].interp(rho_poloidal=self.plasma.rho, t=self.plasma.t) + self.plasma.pressure_fast_parallel.values = Pblon.values + Pbper = code_data["pbper"].interp(rho_poloidal=self.plasma.rho, t=self.plasma.t) + self.plasma.pressure_fast_perpendicular.values = Pbper.values + self.astra_data = code_data + + if self.set_ts_profiles: + overwritten_params = [param for param in self.param_names if any(xs in param for xs in ["Te", "Ne"])] + if any(overwritten_params): + raise ValueError(f"Te/Ne set by TS but then the following params overwritten: {overwritten_params}") + Te = code_data["te"].interp(rho_poloidal=self.plasma.rho, t=self.plasma.time_to_calculate) * 1e3 + self.plasma.Te_prof = lambda: Te.values + Ne = code_data["ne"].interp(rho_poloidal=self.plasma.rho, t=self.plasma.time_to_calculate) * 1e19 + self.plasma.Ne_prof = lambda: Ne.values + + if self.set_all_profiles: + overwritten_params = [param for param in self.param_names if any(xs in param for xs in ["Te", "Ti", "Ne", "Nimp"])] + if any(overwritten_params): + raise ValueError(f"Te/Ne set by TS but then the following params overwritten: {overwritten_params}") + Te = code_data["te"].interp(rho_poloidal=self.plasma.rho, t=self.plasma.time_to_calculate) * 1e3 + self.plasma.Te_prof = lambda: Te.values + Ne = code_data["ne"].interp(rho_poloidal=self.plasma.rho, t=self.plasma.time_to_calculate) * 1e19 + self.plasma.Ne_prof = lambda: Ne.values + Ti = code_data["ti"].interp(rho_poloidal=self.plasma.rho, t=self.plasma.time_to_calculate) * 1e3 + self.plasma.Ti_prof = lambda: Ti.values + Nimp = code_data["niz1"].interp(rho_poloidal=self.plasma.rho, t=self.plasma.time_to_calculate) * 1e19 + self.plasma.Nimp_prof = lambda: Nimp.values + def setup_opt_data(self, phantoms=False, **kwargs): if not hasattr(self, "plasma"): raise ValueError("Missing plasma object required for setup_opt_data") @@ -283,6 +395,17 @@ def _phantom_data(self, noise=False, noise_factor=0.1, **kwargs): .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) ) opt_data["cxff_pi.ti"] = cxrs_data.where(cxrs_data.channel == 2, drop=True) + + if "ts" in self.diagnostics: + _ts_data = self.models["ts"]() + ts_data = {key: _ts_data[key].expand_dims(dim={"t": [self.plasma.time_to_calculate]}) for key in ["te", "ne"]} + opt_data["ts.te"] = ts_data["te"] + opt_data["ts.ne"] = ts_data["ne"] + opt_data["ts.te"]["error"] = opt_data["ts.te"] / opt_data["ts.te"] * ( + 0.10 * opt_data["ts.te"].max(dim="channel")) + opt_data["ts.ne"]["error"] = opt_data["ts.ne"] / opt_data["ts.ne"] * ( + 0.10 * opt_data["ts.ne"].max(dim="channel")) + if "efit" in self.diagnostics: opt_data["efit.wp"] = ( self.models["efit"]() @@ -290,7 +413,10 @@ def _phantom_data(self, noise=False, noise_factor=0.1, **kwargs): .expand_dims(dim={"t": [self.plasma.time_to_calculate]}) ) + # TODO: add chers + if noise: + #TODO: add TS opt_data["smmh1.ne"] = opt_data["smmh1.ne"] + opt_data[ "smmh1.ne" ].max().values * np.random.normal(0, noise_factor, None) @@ -311,20 +437,21 @@ def _phantom_data(self, noise=False, noise_factor=0.1, **kwargs): 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 + opt_data["xrcs.spectra"].wavelength ) 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( + self.model_kwargs["xrcs.background"] = background.mean( dim="wavelength" ).sel(t=self.plasma.time_to_calculate) opt_data["xrcs.spectra"]["error"] = np.sqrt( @@ -333,8 +460,35 @@ def _exp_data(self, **kwargs): 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 + opt_data["cxff_pi"]["ti"] != 0, + ) + + opt_data["cxff_pi"]["ti"] = opt_data["cxff_pi"]["ti"].where( + opt_data["cxff_pi"]["ti"].channel > 3, + ) + opt_data["cxff_pi"]["ti"] = opt_data["cxff_pi"]["ti"].where( + opt_data["cxff_pi"]["ti"].channel < 6, + ) + + if "cxff_tws_c" in self.diagnostics: + opt_data["cxff_tws_c"]["ti"] = opt_data["cxff_tws_c"]["ti"].where( + opt_data["cxff_tws_c"]["ti"] != 0, + ) + opt_data["cxff_tws_c"]["ti"] = opt_data["cxff_tws_c"]["ti"].where( + opt_data["cxff_tws_c"]["ti"].channel < 2, ) + + if "ts" in self.diagnostics: + # TODO: fix error, for now flat error + opt_data["ts.te"] = opt_data["ts.te"].where(opt_data["ts.te"].channel >19, drop=True) + opt_data["ts.ne"] = opt_data["ts.ne"].where(opt_data["ts.ne"].channel > 19, drop=True) + + opt_data["ts.te"]["error"] = opt_data["ts.te"].max(dim="channel") * 0.05 + opt_data["ts.ne"]["error"] = opt_data["ts.ne"].max(dim="channel") * 0.05 + + if self.astra_wp: + opt_data["efit.wp"] = self.astra_data["wth"] + return opt_data def setup_optimiser( @@ -342,12 +496,13 @@ def setup_optimiser( model_kwargs, nwalkers=50, burn_frac=0.10, - sample_high_density=False, + sample_method="random", + **kwargs, ): self.model_kwargs = model_kwargs self.nwalkers = nwalkers self.burn_frac = burn_frac - self.sample_high_density = sample_high_density + self.sample_method = sample_method self.bayesmodel = BayesModels( plasma=self.plasma, @@ -367,21 +522,106 @@ def setup_optimiser( 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: + + + def _sample_start_points(self, sample_method: str = "random", nsamples=100, **kwargs): + if sample_method == "high_density": start_points = self.bayesmodel.sample_from_high_density_region( - self.param_names, self.sampler, self.nwalkers, nsamples=100 + self.param_names, self.sampler, self.nwalkers, nsamples=nsamples + ) + + elif sample_method == "ga": + self.start_points = self._sample_start_points(sample_method="random", **kwargs) + samples_in_weird_format = self.ga_opt(**kwargs) + sample_start_points = np.array([idx[1] for idx in samples_in_weird_format]) + + start_points = np.random.normal( + np.mean(sample_start_points, axis=0), + np.std(sample_start_points, axis=0), + size=(self.nwalkers, sample_start_points.shape[1]), + ) + + + elif sample_method == "random": + start_points = self.bayesmodel.sample_from_priors( + self.param_names, size=self.nwalkers ) + else: + print(f"Sample method: {sample_method} not recognised, Defaulting to random sampling") start_points = self.bayesmodel.sample_from_priors( self.param_names, size=self.nwalkers ) return start_points + + def ga_opt(self, num_gens=30, popsize=50, sols_to_return=5, mutation_probability=None, **kwargs) -> list(tuple((float, []))): + """Runs the GA optimization, and returns a number of the best solutions. Uses + a population convergence stopping criteria: fitness does not improve in 3 successive generations, we stop. + + Args: + num_gens (int, optional): Maximum number of generations to run. Defaults to 30. + popsize (int, optional): Population size. Defaults to 50. + sols_to_return (int, optional): How many of the best solutions the function shall return. Defaults to 5. + + Returns: + list(tuple(float, np.Array(float))): list of tuples, where first element is fitness, second np.array of the parameters. + """ + + import pygad + import time + + # Packaged evaluation function + def idiot_proof(ga_instance, x, sol_idx): + res, _ = self.bayesmodel.ln_posterior(dict(zip(self.param_names, x))) + # print(-res) + return float(res) + + print(f"Running GA for a maximum of {num_gens} generations of {popsize} individuals each.") + + # Initialize the GA instance + ga_instance = pygad.GA(num_generations=num_gens, + num_parents_mating=20, + sol_per_pop=popsize, + num_genes=len(self.start_points[0]), + fitness_func=idiot_proof, + initial_population=self.start_points, + save_best_solutions=True, + stop_criteria="saturate_5", + mutation_probability=mutation_probability) + + st = time.time() + # Execute + ga_instance.run() + + print( + f"Time ran: {time.time() - st:.2f} seconds. Ran total of {ga_instance.generations_completed} generations.") + + # Saves the fitness evolution plot + # figure = ga_instance.plot_fitness() + # figure.savefig(f'GA_plot.png', dpi=300) + + # Organizing all the non-inf individuals from the last generation + feasible_indices = [i for i in range(len(ga_instance.last_generation_fitness)) if + ga_instance.last_generation_fitness[i] != -np.inf] + feasible_individuals = [ga_instance.population[i] for i in feasible_indices] + feasible_fitnesses = [ga_instance.last_generation_fitness[i] for i in feasible_indices] + + # for i, item in enumerate(feasible_individuals_with_keywords): + # item["fitness"]=feasible_fitnesses[i + # feasible_individuals_with_keywords=sorted(feasible_individuals_with_keywords,key= lambda d:d['fitness'],reverse=True) + # feasible_individuals_and_fitnesses=[tuple(feasible_fitnesses[i],feasible_individuals[i]) for i in len(feasible_individuals)] + + # Combining the last individuals to a collection and sorting + feasible_individuals_and_fitnesses = [] + for i in range(len(feasible_fitnesses)): + feasible_individuals_and_fitnesses.append(tuple((feasible_fitnesses[i], feasible_individuals[i]))) + feasible_individuals_and_fitnesses = sorted(feasible_individuals_and_fitnesses, key=lambda x: x[0], + reverse=True) + + return feasible_individuals_and_fitnesses[:sols_to_return] + def __call__( self, filepath="./results/test/", @@ -393,20 +633,67 @@ def __call__( burn_frac=0.10, **kwargs, ): + + self.iterations = iterations + self.burn_frac = burn_frac + 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", + mode="EDIT", ) - self.run_sampler(iterations=iterations, burn_frac=burn_frac) - self.save_pickle(filepath=filepath) + self.result = self._build_inputs_dict() + results = [] + + if not self.tsample.shape: + self.tsample = np.array([self.tsample]) + + self.plasma.time_to_calculate = self.tsample[0] + self.start_points = self._sample_start_points( + sample_method=self.sample_method, **kwargs + ) + + for t in self.tsample: + self.plasma.time_to_calculate = t + print(f"Time: {t}") + self.run_sampler(iterations=iterations, burn_frac=burn_frac) + _result = self._build_result_dict() + results.append(_result) + + self.start_points = self.sampler.get_chain()[-1,:,:] + self.sampler.reset() + + _result = dict(_result, ** self.result) - if plot: # currently requires result with DataArrays - plot_bayes_result(self.result, filepath) + self.save_pickle(_result, filepath=filepath, ) + + if plot: # currently requires result with DataArrays + plot_bayes_result(_result, filepath) + + self.result = dict(self.result, ** results[-1]) + profiles = {} + globals = {} + for key, prof in results[0]["PROFILES"].items(): + if key == "RHO_POLOIDAL": + profiles[key] = results[0]["PROFILES"]["RHO_POLOIDAL"] + elif key == "RHO_TOR": + profiles[key] = results[0]["PROFILES"]["RHO_TOR"] + else: + _profs = [result["PROFILES"][key] for result in results] + profiles[key] = xr.concat(_profs, self.tsample) + + + for key, prof in results[0]["GLOBAL"].items(): + _glob = [result["GLOBAL"][key] for result in results] + globals[key] = xr.concat(_glob, self.tsample) + + result = {"PROFILES":profiles, "GLOBAL":globals} + + self.result = dict(self.result, **result,) self.result = self.dict_of_dataarray_to_numpy(self.result) if mds_write: @@ -418,28 +705,36 @@ def __call__( if __name__ == "__main__": run = BayesWorkflowExample( pulse=None, - diagnostics=["xrcs", "efit", "smmh1", "cxff_pi"], + diagnostics=[ + "xrcs", + "efit", + "smmh1", + "cxff_pi", + "ts", + ], param_names=OPTIMISED_PARAMS, opt_quantity=OPTIMISED_QUANTITY, priors=DEFAULT_PRIORS, profile_params=DEFAULT_PROFILE_PARAMS, phantoms=True, + fast_particles=False, tstart=0.02, tend=0.10, dt=0.005, ) run.setup_plasma( - tsample=0.060, + tsample=0.05, ) run.setup_opt_data(phantoms=run.phantoms) - run.setup_optimiser(nwalkers=50, sample_high_density=True, model_kwargs={}) + run.setup_optimiser(nwalkers=50, sample_method="high_density", model_kwargs=run.model_kwargs, nsamples=100) + # run.setup_optimiser(nwalkers=50, sample_method="ga", model_kwargs=run.model_kwargs, num_gens=50, popsize=100, sols_to_return=3, mutation_probability=None) results = run( - filepath="./results/test/", - pulse_to_write=23000101, + filepath=f"./results/test/", + pulse_to_write=25000000, run="RUN01", mds_write=True, plot=True, - burn_frac=0.10, - iterations=100, + burn_frac=0.2, + iterations=500, ) diff --git a/indica/writers/bda_tree.py b/indica/writers/bda_tree.py index 5a5878a8..293b3741 100644 --- a/indica/writers/bda_tree.py +++ b/indica/writers/bda_tree.py @@ -9,8 +9,8 @@ def bda(): nodes = { - "TIME": ("NUMERIC", "time vector, s"), - "TIME_OPT": ("NUMERIC", "time of optimisation, s"), + "TIME": ("NUMERIC", "time vector of optimisation, s"), + "TIME_BINS": ("NUMERIC", "time vector used for binning, s"), "INPUT": { "BURN_FRAC": ("NUMERIC", "Burn in fraction for chains"), "ITER": ("NUMERIC", "Iterations of optimiser"), @@ -104,12 +104,20 @@ def bda(): "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"), + "TI0Z1_ERR": ("SIGNAL", "Central ion temperature of impurity Z1 error, eV"), + "TI0Z2_ERR": ("SIGNAL", "Central ion temperature of impurity Z2 error, eV"), + "TI0Z3_ERR": ("SIGNAL", "Central ion temperature of impurity Z3 error, eV"), + "NI0Z1_ERR": ("SIGNAL", "Central density of impurity Z1 error, m^-3"), + "NI0Z2_ERR": ("SIGNAL", "Central density of impurity Z2 error, m^-3"), + "NI0Z3_ERR": ("SIGNAL", "Central density of impurity Z3 error, m^-3"), + "WP": ("SIGNAL", "Stored energy, J"), + "WTH": ("SIGNAL", "Thermal component of stored energy, J"), + "PTOT": ("SIGNAL", "Total pressure, Pa"), + "PTH": ("SIGNAL", "Thermal pressure, Pa"), + "WP_ERR": ("SIGNAL", "Stored energy error, J"), + "WTH_ERR": ("SIGNAL", "Thermal component of stored energy error, J"), + "PTOT_ERR": ("SIGNAL", "Total pressure error, Pa"), + "PTH_ERR": ("SIGNAL", "Thermal pressure error, Pa"), }, "PHANTOMS": { "FLAG": ("TEXT", "True if phantoms used"), @@ -133,6 +141,7 @@ def bda(): "AUTO_CORR": ("NUMERIC", "Auto-correlation (iteration nwalker)"), "POST_SAMPLE": ("NUMERIC", "Posterior probability samples (sample)"), "PRIOR_SAMPLE": ("NUMERIC", "Prior samples"), + "GELMANRUBIN": ("NUMERIC", "Gelmin-Rubin convergence diagnostic"), }, } return nodes @@ -266,4 +275,4 @@ def query_yes_no( if __name__ == "__main__": - create_nodes(pulse_to_write=23000101) + create_nodes(pulse_to_write=23000101, mode="EDIT", run="RUN02")