Michaelgemmell/bayes workflow

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

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,18 +204,18 @@ 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(
             self.spectra,
             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
 

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

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,

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)

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