From 38ef5df90a82dd1b6d84ca9df88b0fe74ac1ab79 Mon Sep 17 00:00:00 2001 From: Marco Sertoli Date: Wed, 19 Jul 2023 23:09:26 +0100 Subject: [PATCH] Refactored methods to estimate Zeff from bremsstrahlung using inversion or bayesian inference. --- indica/workflows/zeff_profile.py | 317 ++++++++++++++++++++++--------- 1 file changed, 225 insertions(+), 92 deletions(-) diff --git a/indica/workflows/zeff_profile.py b/indica/workflows/zeff_profile.py index 44d2cf5f..c0ed999f 100644 --- a/indica/workflows/zeff_profile.py +++ b/indica/workflows/zeff_profile.py @@ -1,131 +1,178 @@ import emcee +import matplotlib.pylab as plt import numpy as np import pandas as pd import xarray as xr +from xarray import DataArray from indica.bayesmodels import BayesModels from indica.bayesmodels import get_uniform -from indica.models.diode_filters import BremsstrahlungDiode -from indica.models.plasma import Plasma -import indica.readers.read_st40 as read_st40 +from indica.equilibrium import Equilibrium +from indica.models.diode_filters import example_run as example_diode +from indica.models.plasma import example_run as example_plasma +from indica.operators import tomo_1D +import indica.physics as ph +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 +PATHNAME = "./plots/" + +MAIN_ION = "h" +IMPURITIES = ("c",) +IMPURITY_CONCENTRATION = (0.03,) +FULL_RUN = False +N_RAD = 10 + +PATHNAME = "./plots/" + +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(1e16, 5e18), + "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) & (x1 < x2 * 5), 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), +} +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, + "Te_prof.y0": 3000, + "Te_prof.peaking": 2, + "Ti_prof.y0": 5000, + "Ti_prof.peaking": 2, +} +PARAM_NAMES = [ + "Nimp_prof.y0", + "Nimp_prof.y1", + "Nimp_prof.peaking", +] + + # TODO: allow conditional prior usage even when only # one param is being optimisied i.e. 1 is constant -def bayesian_inference( +def prepare_data(pulse, plasma, model, phantom_data: bool = True): + if pulse is not None: + st40 = ReadST40(pulse, tstart=plasma.tstart, tend=plasma.tend, dt=plasma.dt) + st40(["pi", "efit"]) + attrs = st40.binned_data["pi"]["spectra"].attrs + ( + st40.binned_data["pi"]["background"], + st40.binned_data["pi"]["brightness"], + ) = model.integrate_spectra(st40.binned_data["pi"]["spectra"]) + st40.binned_data["pi"]["background"].attrs = attrs + st40.binned_data["pi"]["brightness"].attrs = attrs + + plasma.initialize_variables() + plasma.set_equilibrium(Equilibrium(st40.raw_data["efit"])) + + model.set_los_transform(st40.binned_data["pi"]["spectra"].transform) + model.set_plasma(plasma) + + data = st40.binned_data["pi"]["brightness"].sel(t=plasma.time_to_calculate) + + if phantom_data: + data = model()["brightness"] + + return data + + +def run_bayesian_analysis( pulse, phantom_profile_params, iterations, result_path, - burn_in=0, - tstart=0.02, - tend=0.10, + tstart=0.01, + tend=0.1, dt=0.01, - time_to_calculate=0.04, + burn_in=0, + tsample=3, nwalkers=10, + phantom_data: bool = True, ): - plasma = Plasma( + print("Generating plasma") + plasma = example_plasma( tstart=tstart, tend=tend, dt=dt, - main_ion="h", - impurities=("c",), - impurity_concentration=(0.2,), - full_run=False, - n_rad=10, + main_ion=MAIN_ION, + impurities=IMPURITIES, + impurity_concentration=IMPURITY_CONCENTRATION, + full_run=FULL_RUN, + n_rad=N_RAD, ) - plasma.build_atomic_data() - plasma.time_to_calculate = plasma.t.sel(t=time_to_calculate, method="nearest") - - 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.time_to_calculate = plasma.t[tsample] 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"] + print("Generating model") + _, pi_model, bckc = example_diode(plasma=plasma) + pi_model.name = "pi" - # 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) + print("Preparing data") + data = prepare_data(pulse, plasma, pi_model, phantom_data=phantom_data) - 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 + phantom_profiles = { + "electron_density": plasma.electron_density.sel(t=plasma.time_to_calculate), + "electron_temperature": plasma.electron_temperature.sel( + t=plasma.time_to_calculate ), - "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), + "ion_temperature": plasma.ion_temperature.sel( + t=plasma.time_to_calculate, element=IMPURITIES[0] + ), + "impurity_density": plasma.impurity_density.sel( + t=plasma.time_to_calculate, element=IMPURITIES[0] + ), + "zeff": plasma.zeff.sel(t=plasma.time_to_calculate, element=IMPURITIES[0]), } - 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", - ] + flat_data = {} + flat_data["pi.brightness"] = data.expand_dims(dim={"t": [plasma.time_to_calculate]}) + print("Instatiating Bayes model") bm = BayesModels( plasma=plasma, data=flat_data, - diagnostic_models=[pi], + diagnostic_models=[pi_model], quant_to_optimise=[ "pi.brightness", ], - priors=priors, + priors=PRIORS, ) - ndim = param_names.__len__() - start_points = bm.sample_from_priors(param_names, size=nwalkers) + 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, + parameter_names=PARAM_NAMES, moves=move, - kwargs={"moment_analysis": False, "calc_spectra": True}, ) + print("Sampling") autocorr = sample_with_autocorr( sampler, start_points, iterations=iterations, auto_sample=5 ) @@ -141,13 +188,13 @@ def bayesian_inference( } samples = sampler.get_chain(flat=True) - prior_samples = bm.sample_from_priors(param_names, size=int(1e5)) + 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, + "param_names": PARAM_NAMES, "phantom_profiles": phantom_profiles, "plasma": plasma, "autocorr": autocorr, @@ -156,7 +203,7 @@ def bayesian_inference( plot_bayes_result(**result, figheader=result_path) -def run_bayesian_inference(): +def bayes(pulse: int = 10821, nwalkers: int = 50, iterations: int = 200): phantom_profile_params = { "Ne_prof.y0": 5e19, "Ne_prof.wcenter": 0.4, @@ -172,20 +219,106 @@ def run_bayesian_inference(): "Ti_prof.y0": 5000, "Ti_prof.peaking": 2, } - nwalkers = 10 # 50 # - iterations = 50 # 200 # - - pathname = "./plots/" - - bayesian_inference( - 10607, + ff = run_bayesian_analysis( + pulse, phantom_profile_params, iterations, - pathname, + PATHNAME, burn_in=0, nwalkers=nwalkers, ) + return ff + + +def inversion( + pulse, + tstart=0.01, + tend=0.1, + dt=0.01, + reg_level_guess: float = 0.3, + phantom_data: bool = True, +): + print("Generating plasma") + plasma = example_plasma( + tstart=tstart, + tend=tend, + dt=dt, + main_ion=MAIN_ION, + impurities=IMPURITIES, + impurity_concentration=IMPURITY_CONCENTRATION, + full_run=FULL_RUN, + n_rad=N_RAD, + ) + + print("Generating model") + _, pi_model, bckc = example_diode(plasma=plasma) + pi_model.name = "pi" + + print("Preparing data") + data = prepare_data(pulse, plasma, pi_model, phantom_data=phantom_data) + if phantom_data: + emissivity = pi_model.emissivity + else: + emissivity = None + + los_transform = data.transform + equilibrium = los_transform.equilibrium + z = los_transform.z + R = los_transform.R + dl = los_transform.dl + + data_t0 = data.isel(t=0).data + has_data = np.logical_not(np.isnan(data.isel(t=0).data)) & (data_t0 > 0) + rho_equil = equilibrium.rho.interp(t=data.t) + input_dict = dict( + brightness=data.data, + dl=dl, + t=data.t.data, + R=R, + z=z, + rho_equil=dict( + R=rho_equil.R.data, + z=rho_equil.z.data, + t=rho_equil.t.data, + rho=rho_equil.data, + ), + has_data=has_data, + debug=False, + ) + if emissivity is not None: + input_dict["emissivity"] = emissivity + + tomo = tomo_1D.SXR_tomography(input_dict, reg_level_guess=reg_level_guess) + tomo() + + pi_model.los_transform.plot() + tomo.show_reconstruction() + + inverted_emissivity = DataArray( + tomo.emiss, coords=[("t", tomo.tvec), ("rho_poloidal", tomo.rho_grid_centers)] + ) + inverted_error = DataArray( + tomo.emiss_err, + coords=[("t", tomo.tvec), ("rho_poloidal", tomo.rho_grid_centers)], + ) + inverted_emissivity.attrs["error"] = inverted_error + data_tomo = data + # bckc_tomo = DataArray(tomo.backprojection, coords=data_tomo.coords) + + zeff = ph.zeff_bremsstrahlung( + plasma.electron_temperature, + plasma.electron_density, + pi_model.filter_wavelength, + bremsstrahlung=inverted_emissivity, + gaunt_approx="callahan", + ) + + plt.figure() + plasma.zeff.sum("element").sel(t=0.03).plot(label="Phantom") + zeff.sel(t=0.03).plot(marker="o", label="Recalculated") + plt.ylabel("Zeff") + plt.legend() if __name__ == "__main__": - run_bayesian_inference() + bayes()