Merge branch 'st40' into bda

indica/converters/time.py

@@ -329,7 +329,7 @@ def bin_in_time_dt(tstart: float, tend: float, dt: float, data: DataArray) -> Da
     -------
     :
         Array like the input, but binned along the time axis.
-
+    TODO: add possibility of doing 50% overlap of time bins!
     """
     check_bounds_bin(tstart, tend, dt, data)
     tlabels = get_tlabels_dt(tstart, tend, dt)

indica/equilibrium.py

@@ -65,12 +65,29 @@ def __init__(
         self.faxs = equilibrium_data["faxs"]
         self.fbnd = equilibrium_data["fbnd"]
         self.ftor = equilibrium_data["ftor"]
+        self.rmag = equilibrium_data["rmag"]
+        self.rbnd = equilibrium_data["rbnd"]
+        self.zmag = equilibrium_data["zmag"]
+        self.zbnd = equilibrium_data["zbnd"]
+        self.zx = self.zbnd.min("arbitrary_index")
         self.rhotor = np.sqrt(
             (self.ftor - self.ftor.sel(rho_poloidal=0.0))
             / (self.ftor.sel(rho_poloidal=1.0) - self.ftor.sel(rho_poloidal=0.0))
         )
         self.psi = equilibrium_data["psi"]
-        self.rho = np.sqrt((self.psi - self.faxs) / (self.fbnd - self.faxs))
+
+        # Including workaround in case faxs or fbnd had messy data
+        rho: DataArray = np.sqrt((self.psi - self.faxs) / (self.fbnd - self.faxs))
+        if np.any(np.isnan(rho.interp(R=self.rmag, z=self.zmag))):
+            self.faxs = self.psi.interp(R=self.rmag, z=self.zmag).drop(["R", "z"])
+            self.fbnd = self.psi.interp(R=self.rbnd, z=self.zbnd).mean(
+                "arbitrary_index"
+            )
+            rho = np.sqrt((self.psi - self.faxs) / (self.fbnd - self.faxs))
+        if np.any(np.isnan(rho)):
+            rho = xr.where(rho > 0, rho, 0.0)
+
+        self.rho = rho
         self.t = self.rho.t
         if "vjac" in equilibrium_data and "ajac" in equilibrium_data:
             self.psin = equilibrium_data["psin"]
@@ -85,11 +102,6 @@ def __init__(
         if "rmji" and "rmjo" in equilibrium_data:
             self.rmji = equilibrium_data["rmji"]
             self.rmjo = equilibrium_data["rmjo"]
-        self.rmag = equilibrium_data["rmag"]
-        self.rbnd = equilibrium_data["rbnd"]
-        self.zmag = equilibrium_data["zmag"]
-        self.zbnd = equilibrium_data["zbnd"]
-        self.zx = self.zbnd.min("arbitrary_index")
         self.R_offset = R_shift
         self.z_offset = z_shift
 

indica/models/diode_filters.py

@@ -31,7 +31,7 @@ def __init__(
         filter_fwhm: float = 1,  # 1
         filter_type: str = "boxcar",
         etendue: float = 1.0,
-        calibration: float = 2.0e-5,
+        calibration: float = 1,
         instrument_method="get_diode_filters",
         channel_mask: slice = None,  # =slice(18, 28),
     ):

indica/models/thomson_scattering.py

@@ -166,7 +166,8 @@ def example_run(
                 alpha=0.7,
             )
             Ne = bckc["ne"].sel(t=t, method="nearest")
-            plt.scatter(Ne.rho_poloidal, Ne, color=cols_time[i], marker="o", alpha=0.7)
+            rho = Ne.transform.rho.sel(t=t, method="nearest")
+            plt.scatter(rho, Ne, color=cols_time[i], marker="o", alpha=0.7)
         plt.xlabel("Channel")
         plt.ylabel("Measured electron density (m^-3)")
         plt.legend()
@@ -179,7 +180,8 @@ def example_run(
                 alpha=0.7,
             )
             Te = bckc["te"].sel(t=t, method="nearest")
-            plt.scatter(Te.rho_poloidal, Te, color=cols_time[i], marker="o", alpha=0.7)
+            rho = Te.transform.rho.sel(t=t, method="nearest")
+            plt.scatter(rho, Te, color=cols_time[i], marker="o", alpha=0.7)
         plt.xlabel("Channel")
         plt.ylabel("Measured electron temperature (eV)")
         plt.legend()

indica/operators/spline_fit_easy.py

@@ -99,63 +99,86 @@ def residuals(yknots):
 
 
 def spline_fit_ts(
-    pulse: int,
-    tstart: float = 0.0,
-    tend: float = 0.2,
-    plot: bool = False,
+    pulse: int = 11314,
+    tstart: float = 0.03,
+    tend: float = 0.1,
+    dt: float = 0.01,
+    quantity: str = "te",
+    R_shift: float = 0.0,
+    knots: list = None,
+    plot: bool = True,
 ):
+    st40 = ReadST40(pulse, tstart=tstart, tend=tend, dt=dt)
+    st40(["ts"], R_shift=R_shift)
+
+    if quantity == "te" and knots is None:
+        knots = [0, 0.3, 0.6, 0.8, 1.1]
+    if quantity == "ne" and knots is None:
+        knots = [0, 0.3, 0.6, 0.8, 0.95, 1.1]
+    data_all = st40.raw_data["ts"][quantity]
+    t = data_all.t
+    transform = data_all.transform
+    transform.convert_to_rho_theta(t=data_all.t)
 
-    ST40 = ReadST40(pulse, tstart=tstart, tend=tend)
-    ST40(["ts"])
-
-    Te_data_all = ST40.binned_data["ts"]["te"]
-    t = ST40.binned_data["ts"]["te"].t
-    transform = ST40.binned_data["ts"]["te"].transform
     R = transform.R
     Rmag = transform.equilibrium.rmag.interp(t=t)
 
     # Fit all available TS data
-    ind = np.full_like(Te_data_all, True)
+    ind = np.full_like(data_all, True)
     rho = xr.where(ind, transform.rho, np.nan)
-    Te_data = xr.where(ind, Te_data_all, np.nan)
-    Te_err = xr.where(ind, Te_data_all.error, np.nan)
-    Te_fit = fit_profile(
-        rho, Te_data, Te_err, knots=[0, 0.3, 0.5, 0.75, 0.95, 1.05], virtual_knots=False
-    )
+    data = xr.where(ind, data_all, np.nan)
+    err = xr.where(ind, data_all.error, np.nan)
+    fit = fit_profile(rho, data, err, knots=knots, virtual_knots=False)
 
     # Use only HFS channels
     ind = R <= Rmag
-    rho = xr.where(ind, transform.rho, np.nan)
-    Te_data = xr.where(ind, Te_data_all, np.nan)
-    Te_err = xr.where(ind, Te_data_all.error, np.nan)
-    Te_fit_hfs = fit_profile(rho, Te_data, Te_err, virtual_knots=True)
+    rho_hfs = xr.where(ind, transform.rho, np.nan)
+    data_hfs = xr.where(ind, data_all, np.nan)
+    err_hfs = xr.where(ind, data_all.error, np.nan)
+    fit_hfs = fit_profile(rho_hfs, data_hfs, err_hfs, knots=knots, virtual_knots=True)
+
+    # Use only LFS channels
+    ind = R >= Rmag
+    rho_lfs = xr.where(ind, transform.rho, np.nan)
+    data_lfs = xr.where(ind, data_all, np.nan)
+    err_lfs = xr.where(ind, data_all.error, np.nan)
+    fit_lfs = fit_profile(rho_lfs, data_lfs, err_lfs, knots=knots, virtual_knots=True)
 
     if plot:
-        for t in Te_data_all.t:
+        for t in data_all.t:
             plt.ioff()
             plt.errorbar(
-                Te_data_all.transform.rho.sel(t=t),
-                Te_data_all.sel(t=t),
-                Te_data_all.error.sel(t=t),
+                rho_hfs.sel(t=t),
+                data_hfs.sel(t=t),
+                err_hfs.sel(t=t),
+                marker="o",
+                label="data HFS",
+                color="blue",
+            )
+            plt.errorbar(
+                rho_lfs.sel(t=t),
+                data_lfs.sel(t=t),
+                err_lfs.sel(t=t),
                 marker="o",
-                label="data",
+                label="data LFS",
+                color="red",
             )
-            Te_fit.sel(t=t).plot(
-                linewidth=5, alpha=0.5, color="orange", label="spline fit all"
+            fit.sel(t=t).plot(
+                linewidth=5, alpha=0.5, color="black", label="spline fit all"
             )
-            Te_fit_hfs.sel(t=t).plot(
-                linewidth=5, alpha=0.5, color="red", label="spline fit HFS"
+            fit_lfs.sel(t=t).plot(
+                linewidth=5, alpha=0.5, color="red", label="spline fit LFS"
+            )
+            fit_hfs.sel(t=t).plot(
+                linewidth=5, alpha=0.5, color="blue", label="spline fit HFS"
             )
             plt.legend()
             plt.show()
 
-    return Te_data_all, Te_fit
+    return data_all, fit
 
 
 if __name__ == "__main__":
-    spline_fit_ts(
-        10619,
-        tstart=0.0,
-        tend=0.2,
-        plot=True,
-    )
+    plt.ioff()
+    spline_fit_ts(11089, quantity="ne")
+    plt.show()