Marcosertoli/satakihino (#263)

* Temporarily stop caching (issues to be solved) * Copied Bremss bayes opt to new file --------- Co-authored-by: Marco Sertoli <[email protected]>
indica-mcf · Jul 17, 2023 · 4dbeddd · 4dbeddd
1 parent 94bfd8d
commit 4dbeddd
Show file tree

Hide file tree

Showing 3 changed files with 238 additions and 54 deletions.
diff --git 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
@@ -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,34 +89,36 @@ 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",
     ]
 
     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)