Refactored methods to estimate Zeff from bremsstrahlung using inversi…

…on or bayesian inference.

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()