diff --git a/indica/converters/time.py b/indica/converters/time.py index cd4ced5f..7c35bd72 100644 --- a/indica/converters/time.py +++ b/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) diff --git a/indica/operators/spline_fit_easy.py b/indica/operators/spline_fit_easy.py index a012f5c0..d6549067 100644 --- a/indica/operators/spline_fit_easy.py +++ b/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() diff --git a/indica/readers/available_quantities.py b/indica/readers/available_quantities.py index bb5d9fb9..133d892e 100644 --- a/indica/readers/available_quantities.py +++ b/indica/readers/available_quantities.py @@ -1,5 +1,6 @@ """ Quantities that can be read with the current abstract reader implementation +TODO: change the tuple to DataArray (long_name, units) - see examples in abstractreader """ from typing import Dict diff --git a/indica/readers/read_st40.py b/indica/readers/read_st40.py index bcd2c577..4589dd00 100644 --- a/indica/readers/read_st40.py +++ b/indica/readers/read_st40.py @@ -45,7 +45,7 @@ "sxrc_xy1": {"brightness": (0, np.inf)}, "sxrc_xy2": {"brightness": (0, np.inf)}, "blom_xy1": {"brightness": (0, np.inf)}, - "ts": {"te": (0, 10.0e3), "ne": (0, 1.0e21)}, + "ts": {"te": (0, np.inf), "ne": (0, np.inf)}, "pi": {"spectra": (0, np.inf)}, "tws_c": {"spectra": (0, np.inf)}, } @@ -91,8 +91,12 @@ def __init__( self.tend = tend self.dt = dt - self.reader = ST40Reader(pulse, tstart - 0.05, tend + 0.05, tree=tree) - self.reader_equil = ST40Reader(pulse, tstart - 0.1, tend + 0.1, tree=tree) + _tend = tend + dt * 2 + _tstart = tstart - dt * 2 + if _tstart < 0: + _tstart = 0.0 + self.reader = ST40Reader(pulse, _tstart, _tend, tree=tree) + self.reader_equil = ST40Reader(pulse, _tstart, _tend, tree=tree) self.equilibrium: Equilibrium self.raw_data: dict = {} diff --git a/indica/readers/st40reader.py b/indica/readers/st40reader.py index 3ba45951..92f567c2 100644 --- a/indica/readers/st40reader.py +++ b/indica/readers/st40reader.py @@ -3,6 +3,7 @@ """ +from copy import deepcopy from typing import Any from typing import Dict from typing import List @@ -96,6 +97,7 @@ class ST40Reader(DataReader): "tws_c": "get_spectrometer", "ts": "get_thomson_scattering", } + # TODO: this will not be necessary once the MDS+ standardisation is complete UIDS_MDS = { "xrcs": "sxr", "princeton": "spectrom", @@ -296,7 +298,8 @@ class ST40Reader(DataReader): }, } - _IMPLEMENTATION_QUANTITIES = { # TODO: these will be different diagnostics!!!!!!! + # TODO: this can be deleted once MDS+ standardisation is complete + _IMPLEMENTATION_QUANTITIES = { "diode_arrays": { # GETTING THE DATA OF THE SXR CAMERA "sxr_camera_1": ("brightness", "total"), "sxr_camera_2": ("brightness", "50_Al_filtered"), @@ -305,6 +308,7 @@ class ST40Reader(DataReader): }, } + # TODO: this can be deleted once MDS+ standardisation is complete _RADIATION_RANGES = { "sxr_camera_1": (1, 20), "sxr_camera_2": (21, 40), @@ -768,6 +772,7 @@ def _get_charge_exchange( z, z_path = self._get_signal(uid, instrument, ":z", revision) R, R_path = self._get_signal(uid, instrument, ":R", revision) + # TODO: temporary fix until geometry sorted (especially pulse if statement..) try: location, location_path = self._get_signal( uid, instrument, ".geometry:location", revision @@ -779,7 +784,6 @@ def _get_charge_exchange( location = np.array([location]) direction = np.array([direction]) - # TODO: temporary fix until geometry sorted if location.shape[0] != x.shape[0]: if self.pulse > 10200: index = np.arange(18, 36) @@ -787,7 +791,6 @@ def _get_charge_exchange( index = np.arange(21, 36) location = location[index] direction = direction[index] - except TreeNNF: location = None direction = None @@ -812,10 +815,8 @@ def _get_charge_exchange( ) except TreeNNF: qval_err = np.full_like(qval, 0.0) - # q_path_err = "" dimensions, _ = self._get_signal_dims(q_path, len(qval.shape)) - results[q + "_records"] = q_path results[q] = qval results[f"{q}_error"] = qval_err @@ -868,13 +869,6 @@ def _get_spectrometer( location = np.array([location]) direction = np.array([direction]) - # if self.pulse > 10200: - # index = np.arange(18, 36) - # else: - # index = np.arange(21, 36) - # location = location[index] - # direction = direction[index] - for q in quantities: qval, q_path = self._get_signal( uid, @@ -892,7 +886,6 @@ def _get_spectrometer( ) except TreeNNF: qval_err = np.full_like(qval, 0.0) - # q_path_err = "" dimensions, _ = self._get_signal_dims(q_path, len(qval.shape)) @@ -953,6 +946,8 @@ def _get_diode_filters( _labels, _ = self._get_signal(uid, instrument, ":label", revision) if type(_labels[0]) == np.bytes_: labels = np.array([label.decode("UTF-8") for label in _labels]) + else: + labels = _labels results["times"] = times results["labels"] = labels @@ -1085,6 +1080,8 @@ def _get_thomson_scattering( revision = results["revision"] times, times_path = self._get_signal(uid, instrument, ":time", revision) + # TODO: hardcoded correction to TS coordinates to be fixed in MDS+ + print("\n Hardcoded correction to TS coordinates to be fixed in MDS+ \n") # location, location_path = self._get_signal( # uid, instrument, ".geometry:location", revision # ) @@ -1095,6 +1092,9 @@ def _get_thomson_scattering( y, y_path = self._get_signal(uid, instrument, ":y", revision) z, z_path = self._get_signal(uid, instrument, ":z", revision) R, R_path = self._get_signal(uid, instrument, ":R", revision) + z = R * 0.0 + x = deepcopy(R) + y = 0 for q in quantities: qval, q_path = self._get_signal( @@ -1170,6 +1170,7 @@ def get_revision_name(self, revision) -> str: """Return string defining RUN## or BEST if revision = 0""" if type(revision) == int: + rev_str = "" if revision < 0: rev_str = "" elif revision == 0: diff --git a/indica/workflows/load_modelling_plasma.py b/indica/workflows/load_modelling_plasma.py index 957577bc..d3167b0b 100644 --- a/indica/workflows/load_modelling_plasma.py +++ b/indica/workflows/load_modelling_plasma.py @@ -11,6 +11,7 @@ from indica.converters.line_of_sight import LineOfSightTransform from indica.converters.transect import TransectCoordinates from indica.equilibrium import Equilibrium +from indica.models.bolometer_camera import Bolometer from indica.models.charge_exchange import ChargeExchange from indica.models.diode_filters import BremsstrahlungDiode from indica.models.helike_spectroscopy import HelikeSpectrometer @@ -29,6 +30,7 @@ CMAP, COLORS = set_plot_colors() DIAGNOSTIC_MODELS = { + "smmh": Interferometry, "smmh1": Interferometry, "nirh1": Interferometry, "xrcs": HelikeSpectrometer, @@ -37,10 +39,23 @@ "cxff_tws_c": ChargeExchange, "cxqf_tws_c": ChargeExchange, "brems": BremsstrahlungDiode, - "sxr_diode_1": SXRcamera, - "sxr_camera_4": SXRcamera, + "sxrc_xy1": Bolometer, + "sxrc_xy2": SXRcamera, + "sxr_spd": SXRcamera, + "blom_xy1": Bolometer, } -INSTRUMENTS = ["smmh1", "nirh1", "xrcs", "sxr_diode_1", "efit", "brems"] +INSTRUMENTS: list = [ + "smmh", + "smmh1", + "nirh1", + "xrcs", + "blom_xy1", + "sxrc_xy1", + "sxrc_xy2", + "efit", + "sxr_spd", +] +REVISIONS: dict = {instr: 0 for instr in INSTRUMENTS} FIG_PATH = f"/home/{getpass.getuser()}/figures/Indica/load_modelling_examples/" plt.ion() @@ -201,9 +216,9 @@ def initialize_diagnostic_models( def example_run( dt: float = 0.01, verbose: bool = True, - example: str = "predictive", + example: str = "alsu_11314", all_runs: bool = False, - plot: bool = False, + plot: bool = True, save_fig: bool = False, fig_style="profiles", alpha: float = 1.0, @@ -216,6 +231,7 @@ def example_run( pulse_code, pulse, equil, + equil_run, code, runs, comment, @@ -237,7 +253,8 @@ def example_run( if pulse is not None: print("Reading ST40 data") st40 = ReadST40(pulse, tstart, tend, dt=dt, tree="st40") - st40(instruments=INSTRUMENTS, map_diagnostics=False) + REVISIONS["efit"] = equil_run + st40(instruments=INSTRUMENTS, map_diagnostics=False, revisions=REVISIONS) if equil != code: equilibrium = {run: Equilibrium(st40.raw_data[equil]) for run in runs} @@ -283,7 +300,10 @@ def example_run( models[instrument].set_plasma(_plasma) - _bckc = models[instrument]() + if instrument == "xrcs": + _bckc = models[instrument](moment_analysis=True) + else: + _bckc = models[instrument]() bckc[run][instrument] = deepcopy(_bckc) if plot or save_fig: @@ -501,11 +521,13 @@ def plot_data_bckc_comparison( run_str += f"-{runs[1]}" fig_path = f"{FIG_PATH}{pulse_code}_{time}_{code}_{run_str}/" - norm = {} - norm["brems"] = True - norm["sxr_camera_4"] = True - norm["sxr_diode_1"] = True - norm["xrcs"] = True + norm: dict = {} + norm["xrcs"] = {} + norm["xrcs"]["spectra"] = True + # norm["brems"] = True + # norm["sxr_camera_4"] = True + # norm["sxrc_xy2"] = {} + # norm["sxrc_xy2"]["brightness"] = True y0 = {} y0["nirh1"] = True y0["smmh1"] = True @@ -518,39 +540,38 @@ def plot_data_bckc_comparison( for run in runs: bckc[run]["efit"] = {"wp": plasma[run].wp} - for instrument in st40.raw_data.keys(): - for quantity in bckc[run][instrument].keys(): - print(instrument) - print(f" {quantity}") - run = runs[0] - + instruments = st40.raw_data.keys() + for instrument in instruments: + quantities = bckc[runs[0]][instrument].keys() + for quantity in quantities: if ( quantity not in st40.binned_data[instrument].keys() or quantity not in st40.raw_data[instrument].keys() ): continue - plt.figure() + print(instrument) + print(f" {quantity}") _raw = st40.raw_data[instrument][quantity] _binned = st40.binned_data[instrument][quantity] - _bckc = bckc[run][instrument][quantity] + _bckc = bckc[runs[0]][instrument][quantity] + tslice = slice(_bckc.t.min().values, _bckc.t.max().values) + str_to_add = "" + tslice_raw = tslice + tslice_binned = tslice + if "error" not in _binned.attrs: _binned.attrs["error"] = xr.full_like(_binned, 0.0) if "stdev" not in _binned.attrs: _binned.attrs["stdev"] = xr.full_like(_binned, 0.0) - err = np.sqrt(_binned.error**2 + _binned.stdev**2) err = xr.where(err / _binned.values < 1.0, err, 0.0) - if len(_bckc.dims) > 1: + if len(_binned.dims) > 1: str_to_add = f" @ {time:.3f} s" tslice_binned = _binned.t.sel(t=time, method="nearest") tslice_raw = _raw.t.sel(t=time, method="nearest") - else: - str_to_add = "" - tslice_raw = tslice - tslice_binned = tslice _raw = _raw.sel(t=tslice_raw) _binned = _binned.sel(t=tslice_binned) @@ -562,6 +583,7 @@ def plot_data_bckc_comparison( _binned -= bgnd _raw -= bgnd + plt.figure() _raw.plot( label="Raw", color=COLORS["raw"], @@ -583,17 +605,14 @@ def plot_data_bckc_comparison( markersize=markersize, ) - label = None + label: str = "Model" for run in runs: - if run == runs[-1]: - label = "Model" - _bckc = bckc[run][instrument][quantity].sel(t=tslice_binned) if instrument in norm.keys(): - _bckc = _bckc * _binned.max() / _bckc.max() - - if label is not None: - label += " (scaled)" + if quantity in norm[instrument].keys(): + _bckc = _bckc / _bckc.max() * _binned.max() + if label is not None: + label += " (scaled)" (_bckc).plot( label=label, @@ -601,91 +620,151 @@ def plot_data_bckc_comparison( linewidth=rcParams["lines.linewidth"] * 1.5, alpha=alpha, ) - set_axis_sci() - plt.title(f"{instrument.upper()} {quantity}" + str_to_add) - if instrument in y0.keys(): - plt.ylim( - 0, - ) + del label - if quantity == "spectra": - # TODO: wavelength axis is sorted from max to min... - plt.xlim(_bckc.wavelength.min(), _bckc.wavelength.max()) + set_axis_sci() + plt.title(f"{instrument.upper()} {quantity}" + str_to_add) + if instrument in y0.keys(): + plt.ylim( + 0, + ) + if quantity == "spectra": + plt.xlim(_bckc.wavelength.min(), _bckc.wavelength.max()) - plt.legend() - save_figure(fig_path, f"{instrument}_{quantity}", save_fig=save_fig) + plt.legend() + save_figure(fig_path, f"{instrument}_{quantity}", save_fig=save_fig) def example_params(example: str, all_runs: bool = False): - comment: str - pulse_code: int - pulse: int - equil: str - code: str - tstart: float - tend: float - tplot: float - runs = [f"RUN{run}" for run in (500 + np.arange(61, 77))] - - if example == "predictive": - comment = "Tests using fixed-boundary predictive ASTRA" - pulse_code = 13110009 - pulse = 10009 - equil = "astra" - code = "astra" - if not all_runs: - runs = ["RUN2621"] - tstart = 0.02 - tend = 0.08 - tplot = 0.06 - elif example == "interpretative_10009": - comment = "interpretative ASTRA using HDA/EFIT" - pulse_code = 13110009 - pulse = 10009 - equil = "efit" - code = "astra" - if not all_runs: - runs = ["RUN564"] - tstart = 0.03 - tend = 0.1 - tplot = 0.06 - elif example == "interpretative_9850": - comment = "ASTRA interpretative using HDA/EFIT" - pulse = 9850 - pulse_code = 13109850 - equil = "efit" - code = "astra" - if not all_runs: - runs = ["RUN564"] - tstart = 0.02 - tend = 0.1 - tplot = 0.08 - elif example == "interpretative_9229": - comment = "ASTRA interpretative using HDA/EFIT" - pulse = 9229 - pulse_code = 13109229 - equil = "efit" - code = "astra" - if not all_runs: - runs = ["RUN572"] - tstart = 0.03 - tend = 0.11 - tplot = 0.06 - elif example == "diverted": - comment = "predictive ASTRA using for diverted scenario" - pulse_code = 13000040 - pulse = 10009 - equil = "astra" - code = "astra" - if not all_runs: - runs = ["RUN292"] - tstart = 0.03 - tend = 0.11 - tplot = 0.1 - else: - raise ValueError(f"No parameters for example {example}") - - return pulse_code, pulse, equil, code, runs, comment, tstart, tend, tplot + runs_all: list = [f"RUN{run}" for run in (500 + np.arange(61, 77))] + + params = { + "predictive": dict( + comment="Tests using fixed-boundary predictive ASTRA", + pulse_code=13110009, + pulse=10009, + equil="astra", + equil_run=0, + code="astra", + runs=["RUN2621"], + tstart=0.02, + tend=0.08, + tplot=0.06, + ), + "interpretative_10009": dict( + comment="interpretative ASTRA using HDA/EFIT", + pulse_code=13110009, + pulse=10009, + equil="efit", + equil_run=0, + code="astra", + runs=["RUN564"], + tstart=0.03, + tend=0.1, + tplot=0.06, + ), + "interpretative_9850": dict( + comment="ASTRA interpretative using HDA/EFIT", + pulse=9850, + pulse_code=13109850, + equil="efit", + equil_run=0, + code="astra", + runs=["RUN564"], + tstart=0.02, + tend=0.1, + tplot=0.08, + ), + "interpretative_9229": dict( + comment="ASTRA interpretative using HDA/EFIT", + pulse=9229, + pulse_code=13109229, + equil="efit", + equil_run=0, + code="astra", + runs=["RUN572"], + tstart=0.03, + tend=0.11, + tplot=0.06, + ), + "diverted": dict( + comment="predictive ASTRA using for diverted scenario", + pulse_code=13000040, + pulse=10009, + equil="astra", + equil_run=0, + code="astra", + runs=["RUN292"], + tstart=0.03, + tend=0.11, + tplot=0.1, + ), + "aleksei_11228": dict( + comment="ASTRA using TS and invented Ti shapes", + pulse=11228, + pulse_code=13011228, + equil="efit", # "astra" + equil_run=1, + code="astra", + runs=["RUN610", "RUN611", "RUN612"], + tstart=0.03, + tend=0.11, + tplot=0.08, + ), + "alsu_11312": dict( + comment="ASTRA using TS and peaked Ti scaled to CXRS", + pulse=11312, + pulse_code=33011312, + equil="astra", + equil_run=0, + code="astra", + runs=["RUN21"], + tstart=0.065, + tend=0.095, + tplot=0.075, + ), + "alsu_11314": dict( + comment="ASTRA using TS and peaked Ti scaled to CXRS", + pulse=11314, + pulse_code=33011314, + equil="astra", + equil_run=0, + code="astra", + runs=["RUN12"], + tstart=0.065, + tend=0.095, + tplot=0.075, + ), + "alsu_11317": dict( + comment="ASTRA using TS and peaked Ti scaled to CXRS", + pulse=11317, + pulse_code=33011317, + equil="astra", + equil_run=0, + code="astra", + runs=["RUN9"], + tstart=0.065, + tend=0.095, + tplot=0.075, + ), + } + + _params = params[example] + if all_runs: + _params["runs"] = runs_all + + return ( + _params["pulse_code"], + _params["pulse"], + _params["equil"], + _params["equil_run"], + _params["code"], + _params["runs"], + _params["comment"], + _params["tstart"], + _params["tend"], + _params["tplot"], + ) # def add_gacode_data( @@ -723,4 +802,6 @@ def example_params(example: str, all_runs: bool = False): if __name__ == "__main__": + plt.ioff() example_run() + plt.show()