Merged main in.

indica/converters/transect.py

@@ -69,35 +69,17 @@ def __init__(
         self._machine_dims = machine_dimensions
         self.x = xr.DataArray(
             x_positions,
-            coords={
-                self.x1_name: self.x1.data
-                if isinstance(self.x1, xr.DataArray)
-                else self.x1
-            },
-            dims=self.x1_name,
+            coords={self.x1_name: self.x1.data},
         )
         self.y = xr.DataArray(
             y_positions,
-            coords={
-                self.x1_name: self.x1.data
-                if isinstance(self.x1, xr.DataArray)
-                else self.x1
-            },
-            dims=self.x1.name,
+            coords={self.x1_name: self.x1.data},
         )
         self.z = xr.DataArray(
             z_positions,
-            coords={
-                self.x1_name: self.x1.data
-                if isinstance(self.x1, xr.DataArray)
-                else self.x1
-            },
-            dims=self.x1_name,
+            coords={self.x1_name: self.x1.data},
         )
 
-        # self.x: DataArray = DataArray(x_positions, coords=[(self.x1_name, self.x1)])
-        # self.y: DataArray = DataArray(y_positions, coords=[(self.x1_name, self.x1)])
-        # self.z: DataArray = DataArray(z_positions, coords=[(self.x1_name, self.x1)])
         self.R: DataArray = np.sqrt(self.x**2 + self.y**2)
         self.rho: DataArray
         self.theta: DataArray

indica/datatypes.py

@@ -29,6 +29,9 @@
     "density_integrated": "1/$m^2$",
     "brightness": "$W/m^2$",
     "emissivity": "$W/m^3$",
+    "intensity": "$count/s$",
+    "radiance": "W/m^2/steradian$",
+    "spectral_radiance": "W/m^2/steradian/nm$",
     "velocity": "m/s",
     "magnetic_flux": r"$Wb/2\pi$",
     "f": "Wb m",
@@ -88,6 +91,9 @@
         "Centrifugal asymmetry multiplier",
         "none",
     ),
+    "intensity": ("Intensity", "intensity"),
+    "radiance": ("Radiance", "radiance"),
+    "spectral_radiance": ("Spectral Radiance", "spectral_radiance"),
     "emissivity": ("Emissivity", "emissivity"),
     "total_radiated_power_emission": ("$P_{rad, tot}$", "emissivity"),
     "sxr_radiated_power_emission": ("$P_{rad, sxr}$", "emissivity"),

indica/defaults/load_defaults.py

@@ -0,0 +1,48 @@
+from pathlib import Path
+import pickle
+
+
+PROJECT_PATH = Path(__file__).parent.parent
+DEFAULTS_PATH = f"{PROJECT_PATH}/defaults/"
+
+
+def get_filename_default_objects(machine: str):
+    _files = {}
+    _files["geometry"] = (
+        DEFAULTS_PATH + f"{machine}_default_geometry_transform_objects.pkl"
+    )
+    _files["equilibrium"] = DEFAULTS_PATH + f"{machine}_default_equilibrium_object.pkl"
+    _files["plasma"] = DEFAULTS_PATH + f"{machine}_default_plasma_object.pkl"
+
+    return _files
+
+
+def load_default_objects(machine: str, identifier: str = "geometry"):
+    """
+    Load default objects from local pickle files
+
+    Parameters
+    ----------
+    machine - e.g. "st40"
+    identifier - "geometry" or "equilibrium" or "plasma"
+    """
+    _file = get_filename_default_objects(machine)[identifier]
+
+    try:
+        return pickle.load(open(_file, "rb"))
+    except FileNotFoundError:
+        to_print = f"""
+
+************************************************************
+The following file does not exist:
+{_file}
+
+Create your defaults file:
+    python indica/defaults/read_write_defaults.py
+
+or, to choose specific kwargs (e.g. machine or pulse):
+    save_default_objects("st40", 11419)
+
+************************************************************
+            """
+        raise Exception(to_print)

indica/defaults/read_write_defaults.py → indica/defaults/save_defaults.py

@@ -3,6 +3,7 @@
 
 import numpy as np
 
+from indica.defaults.load_defaults import get_filename_default_objects
 from indica.equilibrium import Equilibrium
 from indica.models import Plasma
 from indica.models.plasma import PlasmaProfiles
@@ -120,47 +121,5 @@ def save_default_objects(
     return plasma
 
 
-def get_filename_default_objects(machine: str):
-    _files = {}
-    _files["geometry"] = (
-        DEFAULTS_PATH + f"{machine}_default_geometry_transform_objects.pkl"
-    )
-    _files["equilibrium"] = DEFAULTS_PATH + f"{machine}_default_equilibrium_object.pkl"
-    _files["plasma"] = DEFAULTS_PATH + f"{machine}_default_plasma_object.pkl"
-
-    return _files
-
-
-def load_default_objects(machine: str, identifier: str = "geometry"):
-    """
-    Load default objects from local pickle files
-
-    Parameters
-    ----------
-    machine - e.g. "st40"
-    identifier - "geometry" or "equilibrium" or "plasma"
-    """
-    _file = get_filename_default_objects(machine)[identifier]
-
-    try:
-        return pickle.load(open(_file, "rb"))
-    except FileNotFoundError:
-        to_print = f"""
-
-************************************************************
-The following file does not exist:
-{_file}
-
-Create your defaults file:
-    python indica/defaults/read_write_defaults.py
-
-or, to choose specific kwargs (e.g. machine or pulse):
-    save_default_objects("st40", 11419)
-
-************************************************************
-            """
-        raise Exception(to_print)
-
-
 if __name__ == "__main__":
-    save_default_objects("st40", 11419)
+    save_default_objects("st40", 11560)

indica/models/helike_spectroscopy.py

@@ -5,8 +5,8 @@
 from xarray import DataArray
 
 from indica.converters import LineOfSightTransform
+from indica.defaults.load_defaults import load_default_objects
 from indica.models.abstractdiagnostic import DiagnosticModel
-from indica.models.plasma import example_plasma
 from indica.numpy_typing import LabeledArray
 import indica.physics as ph
 from indica.readers.available_quantities import AVAILABLE_QUANTITIES
@@ -16,7 +16,6 @@
 from indica.utilities import set_axis_sci
 from indica.utilities import set_plot_rcparams
 
-# TODO: why resonance lines in upper case, others lower?
 LINE_RANGES = {
     "w": slice(0.39489, 0.39494),
     "n3": slice(0.39543, 0.39574),
@@ -328,15 +327,20 @@ def _moment_analysis(self):
 
     def _build_bckc_dictionary(self):
         self.bckc = {}
-        if "spectra" in self.quantities and hasattr(self, "measured_spectra"):
-            self.bckc["spectra"] = self.measured_spectra
+        if hasattr(self, "measured_spectra"):
+            self.bckc["raw_spectra"] = self.measured_spectra
 
         if self.moment_analysis:
             for quantity in self.quantities:
-                if quantity == "spectra":
+                datatype = self.quantities[quantity]
+
+                if quantity in [
+                    "raw_spectra",
+                    "spectra",
+                    "background",
+                ]:
                     continue
 
-                datatype = self.quantities[quantity]
                 line = str(quantity.split("_")[1])
                 if "int" in quantity and line in self.measured_intensity.keys():
                     self.bckc[quantity] = self.measured_intensity[line]
@@ -475,18 +479,18 @@ def plot(self):
         plt.figure()
         channels = self.los_transform.x1
         cols_time = cm.gnuplot2(np.linspace(0.1, 0.75, len(self.t), dtype=float))
-        if "spectra" in self.bckc.keys():
-            spectra = self.bckc["spectra"]
-            if "channel" in spectra.dims:
-                spectra = spectra.sel(channel=np.median(channels))
+        if "raw_spectra" in self.bckc.keys():
+            raw_spectra = self.bckc["raw_spectra"]
+            if "channel" in raw_spectra.dims:
+                raw_spectra = raw_spectra.sel(channel=np.median(channels))
             for i, t in enumerate(np.array(self.t, ndmin=1)):
                 plt.plot(
-                    spectra.wavelength,
-                    spectra.sel(t=t),
+                    raw_spectra.wavelength,
+                    raw_spectra.sel(t=t),
                     color=cols_time[i],
                     label=f"t={t:1.2f} s",
                 )
-            plt.ylabel("Brightness (a.u.)")
+            plt.ylabel("Intensity (count/s.)")
             plt.xlabel("Wavelength (nm)")
             plt.legend()
             set_axis_sci()
@@ -496,7 +500,9 @@ def plot(self):
         for i, t in enumerate(self.t):
             plt.plot(
                 self.plasma.ion_temperature.rho_poloidal,
-                self.plasma.ion_temperature.sel(t=t),
+                self.plasma.ion_temperature.sel(
+                    t=t,
+                ),
                 color=cols_time[i],
             )
             plt.plot(
@@ -581,16 +587,13 @@ def helike_transform_example(nchannels):
     return los_transform
 
 
-def example_run(
-    pulse: int = None, plasma=None, plot=False, moment_analysis: bool = False, **kwargs
-):
-    # TODO: LOS sometime crossing bad EFIT reconstruction
+def example_run(plasma=None, plot=False, moment_analysis: bool = False, **kwargs):
+
     if plasma is None:
-        plasma = example_plasma(
-            pulse=pulse, impurities=("ar",), impurity_concentration=(0.001,), n_rad=10
-        )
-        # plasma.time_to_calculate = plasma.t[3:5]
-        # Create new diagnostic
+        plasma = load_default_objects("st40", "plasma")
+        equilibrium = load_default_objects("st40", "equilibrium")
+        plasma.set_equilibrium(equilibrium)
+
     diagnostic_name = "xrcs"
     los_transform = helike_transform_example(3)
     los_transform.set_equilibrium(plasma.equilibrium)

indica/readers/abstractreader.py

@@ -656,7 +656,7 @@ def get_helike_spectroscopy(
 
         data: dict = {}
         for quantity in quantities:
-            if quantity == "spectra":
+            if quantity in ["spectra", "raw_spectra"]:
                 dims = ["t", "wavelength"]
             else:
                 dims = ["t"]
@@ -1003,7 +1003,7 @@ def assign_dataarray(
                 error = error.sel(t=slice(self._tstart, self._tend))
 
         if include_error:
-            data.attrs["error"] = error
+            data = data.assign_coords(error=(data.dims, error.data))
 
         if transform is not None:
             data.attrs["transform"] = transform

indica/readers/available_quantities.py

@@ -27,15 +27,17 @@
         "zeff": "effective_charge",
     },
     "get_helike_spectroscopy": {
+        "ti_w": "ion_temperature",
+        "ti_z": "ion_temperature",
+        "te_n3w": "electron_temperature",
+        "te_kw": "electron_temperature",
         "int_w": "line_intensity",
         "int_k": "line_intensity",
         "int_tot": "line_intensity",
         "int_n3": "line_intensity",
-        "te_kw": "electron_temperature",
-        "te_n3w": "electron_temperature",
-        "ti_w": "ion_temperature",
-        "ti_z": "ion_temperature",
-        "spectra": "spectra",
+        "raw_spectra": "intensity",
+        "spectra": "spectral_radiance",
+        "background": "intensity",
     },
     "get_ppts": {
         "ne_rho": "electron_density",

indica/readers/read_st40.py

@@ -16,9 +16,9 @@
 
 INSTRUMENTS: list = [
     "efit",
-    "lines",
+    # "lines",
     "nirh1",
-    "smmh1",
+    "smmh",
     "smmh",
     "cxff_pi",
     "cxff_tws_c",
@@ -39,7 +39,14 @@
     "cxff_pi": {"ti": (0, np.inf), "vtor": (0, np.inf)},
     "cxff_tws_c": {"ti": (0, np.inf), "vtor": (0, np.inf)},
     "cxqf_tws_c": {"ti": (0, np.inf), "vtor": (0, np.inf)},
-    "xrcs": {"ti_w": (0, np.inf), "te_kw": (0, np.inf), "te_n3w": (0, np.inf)},
+    "xrcs": {
+        "ti_w": (0, np.inf),
+        "ti_z": (0, np.inf),
+        "te_kw": (0, np.inf),
+        "te_n3w": (0, np.inf),
+        "spectra": (0, np.inf),
+        "raw_spectra": (0, np.inf),
+    },
     "brems": {"brightness": (0, np.inf)},
     "halpha": {"brightness": (0, np.inf)},
     "sxr_spd": {"brightness": (0, np.inf)},
@@ -66,6 +73,7 @@
     "cxff_tws_c": {"ti": ("channel", (0, np.inf)), "vtor": ("channel", (0, np.inf))},
     "xrcs": {
         "spectra": ("wavelength", (0.0, np.inf)),
+        "raw_spectra": ("wavelength", (0.0, np.inf)),
     },
     "ts": {"te": ("channel", (0, np.inf)), "ne": ("channel", (0, np.inf))},
 }
@@ -300,17 +308,17 @@ def __call__(
         self.get_equilibrium(R_shift=R_shift, revision=revisions["efit"])
         for instrument in instruments:
             print(f"Reading {instrument}")
-            try:
-                self.get_raw_data(
-                    "",
-                    instrument,
-                    revisions[instrument],
-                    set_equilibrium=set_equilibrium,
-                )
-            except Exception as e:
-                print(f"error reading: {instrument} \nException: {e}")
-                if debug:
-                    raise e
+            # try:
+            self.get_raw_data(
+                "",
+                instrument,
+                revisions[instrument],
+                set_equilibrium=set_equilibrium,
+            )
+            # except Exception as e:
+            #     print(f"error reading: {instrument} \nException: {e}")
+            #     if debug:
+            #         raise e
 
         if raw_only:
             return
@@ -361,6 +369,14 @@ def bin_data_in_time(
             if "t" in data_quant.coords:
                 data_quant = convert_in_time_dt(tstart, tend, dt, data_quant)
 
+            # Using groupedby_bins always removes error from coords so adding it back
+            if "error" in raw_data[instr][quant].coords:
+                error = convert_in_time_dt(
+                    tstart, tend, dt, raw_data[instr][quant].error
+                )
+                data_quant = data_quant.assign_coords(
+                    error=(raw_data[instr][quant].dims, error.data)
+                )
             binned_quantities[quant] = data_quant
         binned_data[instr] = binned_quantities
     return binned_data
@@ -408,7 +424,7 @@ def coord_condition(data: DataArray, coord_info: tuple):
     condition = (data.coords[coord_name] >= coord_slice[0]) * (
         data.coords[coord_name] < coord_slice[1]
     )
-    filtered_data = xr.where(condition, data, np.nan)
+    filtered_data = data.where(condition, np.nan)
     filtered_data.attrs = data.attrs
     return filtered_data