From f3093f7048b5a46524de7c960f9b7ad746ea8d0b Mon Sep 17 00:00:00 2001 From: Marco Sertoli Date: Fri, 27 Oct 2023 14:56:16 +0100 Subject: [PATCH 01/17] Added development directory to hold files still under development but potentially useful to the user group. --- indica/development/bolometer_vos.py | 164 ++++++++++++++++++++++++++++ 1 file changed, 164 insertions(+) create mode 100644 indica/development/bolometer_vos.py diff --git a/indica/development/bolometer_vos.py b/indica/development/bolometer_vos.py new file mode 100644 index 00000000..5058bab3 --- /dev/null +++ b/indica/development/bolometer_vos.py @@ -0,0 +1,164 @@ +from copy import deepcopy +import getpass + +from matplotlib import rcParams +import matplotlib.pylab as plt +import numpy as np + +from indica.converters.line_of_sight import LineOfSightTransform +from indica.models.bolometer_camera import Bolometer +from indica.models.plasma import example_run as example_plasma +from indica.models.plasma import Plasma +from indica.operators import tomo_1D +from indica.readers.read_st40 import ReadST40 +from indica.utilities import save_figure +from indica.utilities import set_axis_sci +from indica.utilities import set_plot_colors +from indica.utilities import set_plot_rcparams + +FIG_PATH = f"/home/{getpass.getuser()}/figures/Indica/bolometer_vos/" + +CMAP, COLORS = set_plot_colors() +set_plot_rcparams("profiles") +rcParams.update({"font.size": 16}) + + +def bolo_xy_p23( + pulse: int = 11336, + instrument: str = "sxrc_xy1", + plasma: Plasma = None, + st40: ReadST40 = None, + save_fig: bool = False, +): + + if plasma is None: + plasma = example_plasma(pulse=pulse, calc_power_loss=True) + + if st40 is None: + st40 = ReadST40(pulse) + st40(instruments=[instrument]) + quantity = list(st40.raw_data[instrument])[0] + los_transform = st40.raw_data[instrument][quantity].transform + los_transform.set_equilibrium(plasma.equilibrium, force=True) + origin = los_transform.origin + direction = los_transform.direction + + model = Bolometer( + instrument, + ) + + # Nominal LOS geometry + model.set_los_transform(los_transform) + model.set_plasma(plasma) + + # Shifted LOS geometry + los_transform_plus = LineOfSightTransform( + origin_x=origin[:, 0] + 0.01, + origin_y=origin[:, 1], + origin_z=origin[:, 2], + direction_x=direction[:, 0], + direction_y=direction[:, 1], + direction_z=direction[:, 2], + machine_dimensions=los_transform._machine_dims, + ) + los_transform_plus.set_equilibrium(plasma.equilibrium) + model_plus = deepcopy(model) + model_plus.set_los_transform(los_transform_plus) + + los_transform_minus = LineOfSightTransform( + origin_x=origin[:, 0] - 0.01, + origin_y=origin[:, 1], + origin_z=origin[:, 2], + direction_x=direction[:, 0], + direction_y=direction[:, 1], + direction_z=direction[:, 2], + machine_dimensions=los_transform._machine_dims, + ) + los_transform_minus.set_equilibrium(plasma.equilibrium) + model_minus = deepcopy(model) + model_minus.set_los_transform(los_transform_minus) + + bckc = model() + bckc_plus = model_plus() + bckc_minus = model_minus() + + cols = model.los_transform.plot(orientation="xy") + for ch in model.los_transform.x1: + plt.plot( + model_plus.los_transform.x.sel(channel=ch), + model_plus.los_transform.y.sel(channel=ch), + color=cols[ch], + linewidth=2, + alpha=0.4, + ) + plt.plot( + model_minus.los_transform.x.sel(channel=ch), + model_minus.los_transform.y.sel(channel=ch), + color=cols[ch], + linewidth=2, + alpha=0.4, + ) + plt.title("LOS geometry (x,y)") + save_figure(FIG_PATH, "los_geometry", save_fig=save_fig) + + cols_time = CMAP(np.linspace(0.1, 0.75, len(model.t), dtype=float)) + plt.figure() + for i, t in enumerate(model.t.values): + nominal = bckc["brightness"].sel(t=t, method="nearest") + plus = bckc_plus["brightness"].sel(t=t, method="nearest") + minus = bckc_minus["brightness"].sel(t=t, method="nearest") + nominal.plot(label=f"t={t:1.2f} s", color=cols_time[i], alpha=0.5) + plt.fill_between(nominal.channel, plus, minus, color=cols_time[i], alpha=0.7) + + set_axis_sci() + plt.xlabel("Channel") + plt.ylabel("[W/$m^2$]") + plt.legend() + plt.title("Measured brightness") + save_figure(FIG_PATH, "bckc_brightness", save_fig=save_fig) + + # Local emissivity profiles + plt.figure() + for i, t in enumerate(model.t.values): + model.emissivity.sel(t=t).plot( + label=f"t={t:1.2f} s", color=cols_time[i], alpha=0.8 + ) + + set_axis_sci() + plt.xlabel("rho") + plt.ylabel("[W/$m^3$]") + plt.legend() + plt.title("Local radiated power") + save_figure(FIG_PATH, "local_radiated_power", save_fig=save_fig) + + data_t0 = bckc["brightness"].isel(t=0).data + has_data = np.logical_not(np.isnan(data_t0)) & (data_t0 >= 1.0e3) + + rho_equil = model.los_transform.equilibrium.rho.interp( + t=bckc["brightness"].t, method="nearest" + ) + input_dict = dict( + brightness=bckc["brightness"].data, + dl=model.los_transform.dl, + t=bckc["brightness"].t.data, + R=model.los_transform.R.data, + z=model.los_transform.z.data, + rho_equil=dict( + R=rho_equil.R.data, + z=rho_equil.z.data, + t=rho_equil.t.data, + rho=rho_equil.data, + ), + emissivity=model.emissivity, + debug=False, + has_data=has_data, + ) + # return input_dict + + tomo = tomo_1D.SXR_tomography(input_dict, reg_level_guess=0.5) + + tomo() + plt.ioff() + tomo.show_reconstruction() + + return plasma, st40 From 2e12a4726e7cdb9d3723fdffd91511f8efeebc1a Mon Sep 17 00:00:00 2001 From: Marco Sertoli Date: Fri, 27 Oct 2023 14:57:04 +0100 Subject: [PATCH 02/17] Modified tomo_1D.py to be aesthetically more in line with the rest of the plotting. --- indica/converters/abstractconverter.py | 2 ++ indica/operators/tomo_1D.py | 21 ++++++++++++--------- 2 files changed, 14 insertions(+), 9 deletions(-) diff --git a/indica/converters/abstractconverter.py b/indica/converters/abstractconverter.py index dd7364f3..f58f9a5d 100644 --- a/indica/converters/abstractconverter.py +++ b/indica/converters/abstractconverter.py @@ -581,6 +581,8 @@ def plot( plt.title(title) save_figure(fig_path, f"{fig_name}{self.name}_rho", save_fig=save_fig) + return cols + def plot_geometry( abscissa: DataArray, diff --git a/indica/operators/tomo_1D.py b/indica/operators/tomo_1D.py index de361b0b..4645daf8 100644 --- a/indica/operators/tomo_1D.py +++ b/indica/operators/tomo_1D.py @@ -499,12 +499,12 @@ def show_reconstruction(self): tomo_var = ax[0].fill_between( r, r * 0, r * 0, alpha=0.5, facecolor="b", edgecolor="None" ) - (expected_emissivity,) = ax[0].plot( - [], [], alpha=0.5, color="r", linestyle="dashed", lw=2 - ) + (expected_emissivity,) = ax[0].plot([], [], color="r", linestyle="dashed", lw=2) (tomo_mean,) = ax[0].plot([], [], lw=2) - errorbar = ax[1].errorbar(0, np.nan, 0, capsize=4, c="g", marker="o", ls="none") + errorbar = ax[1].errorbar( + 0, np.nan, 0, capsize=4, c="r", marker="o", ls="none", alpha=0.5 + ) (retro_inter,) = ax[1].plot([], [], "b-") (retro,) = ax[1].plot([], [], "bx") @@ -522,20 +522,18 @@ def show_reconstruction(self): print(ymax) ax[0].set_xlabel(r"$\rho$") - ax[1].set_xlabel(r"index") + ax[1].set_xlabel(r"Channel") ax[1].set_ylabel("Brightness [kW/m$^2$]") ax[0].set_ylabel("Emissivity [kW/m$^3$]") global cont cvals = np.linspace(0, 1, 20) cont = ax[2].contour( - self.eq["R"], self.eq["z"], self.eq["rho"][0], cvals, colors="k" + self.eq["R"], self.eq["z"], self.eq["rho"][0], cvals, colors="red" ) ax[2].axis("equal") ax[2].plot(self.R.T, self.z.T, "b", zorder=99) - R, z = np.meshgrid(self.eq["R"], self.eq["z"]) - ax[2].plot(R, z, "k,", zorder=99) ax[2].axis( [ 0.15, @@ -589,7 +587,12 @@ def update(reg=None, time=None): it_eq = np.argmin(np.abs(self.eq["t"] - time)) cont = ax[2].contour( - self.eq["R"], self.eq["z"], self.eq["rho"][it_eq], cvals, colors="k" + self.eq["R"], + self.eq["z"], + self.eq["rho"][it_eq], + cvals, + colors=["r"] * len(cvals), + alpha=0.8, ) title.set_text( From a8390a5feda106945d65a37aba9e0fe2b5c34628 Mon Sep 17 00:00:00 2001 From: Marco Sertoli Date: Fri, 27 Oct 2023 15:09:19 +0100 Subject: [PATCH 03/17] Deleted unused plotting methods --- indica/plotters/plot_profiles.py | 231 ------------------------------- 1 file changed, 231 deletions(-) delete mode 100644 indica/plotters/plot_profiles.py diff --git a/indica/plotters/plot_profiles.py b/indica/plotters/plot_profiles.py deleted file mode 100644 index a367ee7c..00000000 --- a/indica/plotters/plot_profiles.py +++ /dev/null @@ -1,231 +0,0 @@ -# import getpass -# -# import matplotlib.pyplot as plt -# import numpy as np -# import xarray as xr -# -# from indica.numpy_typing import ArrayLike -# from indica.operators.gpr_fit import plot_gpr_fit -# from indica.operators.gpr_fit import run_gpr_fit -# from indica.readers.read_st40 import ReadST40 -# from indica.utilities import save_figure -# from indica.utilities import set_plot_colors -# from indica.utilities import set_plot_rcparams -# -# FIG_PATH = f"/home/{getpass.getuser()}/figures/Indica/profiles/" -# CMAP, COLORS = set_plot_colors() -# set_plot_rcparams("profiles") -# -# -# def plot_ts_cxrs_profiles( -# pulse: int = 10605, -# instruments: list = ["ts", "cxff_pi", "cxff_tws_c"], -# tplot: ArrayLike = None, -# save_fig=False, -# max_norm_err: float = 0.05, -# plot_rho: bool = True, -# plot_spectra: bool = True, -# ): -# tstart = 0.02 -# tend = 0.1 -# -# st40 = ReadST40(pulse, tstart, tend) -# st40(instruments=instruments) -# -# _data = st40.raw_data["ts"] -# ts_te = _data["te"].swap_dims({"channel": "R"}) -# ts_rho, _ = _data["te"].transform.convert_to_rho_theta(t=_data["te"].t) -# ts_rho = ts_rho.assign_coords(R=("channel", _data["te"].R)).swap_dims( -# {"channel": "R"} -# ) -# x_bounds = _data["te"].transform._machine_dims[0] -# ts_fit, ts_fit_err = run_gpr_fit( -# ts_te, -# x_bounds=x_bounds, -# y_bounds=(1, 1), -# err_bounds=(0, 0), -# xdim="R", -# ) -# -# if "cxff_pi" in st40.raw_data.keys(): -# chan_slice = slice(3, 5) -# _data = st40.raw_data["cxff_pi"] -# pi_rho, _ = _data["ti"].transform.convert_to_rho_theta(t=_data["ti"].t) -# pi_rho = pi_rho.sel(channel=chan_slice) -# pi_ti = _data["ti"].sel(channel=chan_slice) -# pi_ti_err = _data["ti"].error.sel(channel=chan_slice) -# pi_spectra = _data["spectra"].sel(channel=chan_slice) -# pi_fit = _data["fit"].sel(channel=chan_slice) -# -# swap_dims = {"channel": "R"} -# ti_cond = (pi_ti > 0) * (pi_ti_err / pi_ti <= max_norm_err) -# spectra_cond = pi_ti.expand_dims({"wavelength": pi_spectra.wavelength}) > 0 -# pi_ti = xr.where(ti_cond, pi_ti, np.nan).swap_dims(swap_dims) -# pi_ti_err = xr.where(ti_cond, pi_ti_err, np.nan).swap_dims(swap_dims) -# pi_spectra = xr.where(spectra_cond, pi_spectra, np.nan).swap_dims(swap_dims) -# pi_fit = xr.where(spectra_cond, pi_fit, np.nan).swap_dims(swap_dims) -# if plot_rho: -# pi_rho = pi_rho.assign_coords(R=("channel", pi_ti.R)).swap_dims(swap_dims) -# -# if "cxff_tws_c" in st40.raw_data.keys(): -# chan_slice = slice(0, 2) -# _data = st40.raw_data["cxff_tws_c"] -# tws_rho, _ = _data["ti"].transform.convert_to_rho_theta(t=_data["ti"].t) -# tws_rho = tws_rho.sel(channel=chan_slice) -# tws_ti = _data["ti"].sel(channel=chan_slice) -# tws_ti_err = _data["ti"].error.sel(channel=chan_slice) -# tws_spectra = _data["spectra"].sel(channel=chan_slice) -# tws_fit = _data["fit"].sel(channel=chan_slice) -# -# swap_dims = {"channel": "R"} -# ti_cond = (tws_ti > 0) * (tws_ti_err / tws_ti <= max_norm_err) -# spectra_cond = tws_ti.expand_dims({"wavelength": tws_spectra.wavelength}) > 0 -# tws_ti = xr.where(ti_cond, tws_ti, np.nan).swap_dims(swap_dims) -# tws_ti_err = xr.where(ti_cond, tws_ti_err, np.nan).swap_dims(swap_dims) -# tws_spectra = xr.where(spectra_cond, tws_spectra, np.nan).swap_dims(swap_dims) -# tws_fit = xr.where(spectra_cond, tws_fit, np.nan).swap_dims(swap_dims) -# tws_rho = tws_rho.assign_coords(R=("channel", tws_ti.R)).swap_dims(swap_dims) -# -# plt.ioff() -# fig_name = f"{pulse}_TS_CXRS_profiles" -# if tplot is None: -# tplot = ts_te.t.values -# -# t_pi = -10 -# t_tws = -10 -# tplot = np.array(tplot, ndmin=1) -# for t in tplot: -# t_ts = ts_te.t.sel(t=t, method="nearest").values -# plot_gpr_fit( -# ts_te, -# ts_fit, -# ts_fit_err, -# t_ts, -# ylabel="[eV]", -# xlabel="R [m]", -# title=str(st40.pulse), -# label=f"El. (TS) {t_ts:.3f} s", -# fig_name=f"{fig_name}_vs_R", -# color=COLORS["electron"], -# ) -# if "cxff_pi" in st40.raw_data.keys(): -# t_pi = pi_ti.t.sel(t=t, method="nearest").values -# if "cxff_tws_c" in st40.raw_data.keys(): -# t_tws = tws_ti.t.sel(t=t, method="nearest").values -# if np.abs(t_ts - t_pi) < 0.02: -# plt.errorbar( -# pi_ti.R, -# pi_ti.sel(t=t_pi).values, -# pi_ti_err.sel(t=t_pi).values, -# color=COLORS["ion"], -# ) -# plt.plot( -# pi_ti.R, -# pi_ti.sel(t=t_pi).values, -# marker="s", -# color=COLORS["ion"], -# label=f"Ion (PI) {t_pi:.3f} s", -# ) -# -# if np.abs(t_ts - t_tws) < 0.02: -# plt.errorbar( -# tws_ti.R, -# tws_ti.sel(t=t_tws).values, -# tws_ti_err.sel(t=t_tws).values, -# color=COLORS["ion"], -# ) -# plt.plot( -# tws_ti.R, -# tws_ti.sel(t=t_tws).values, -# marker="^", -# color=COLORS["ion"], -# label=f"Ion (TWS) {t_tws:.3f} s", -# ) -# plt.legend(loc="upper left") -# if save_fig: -# save_figure( -# FIG_PATH, -# f"{fig_name}_vs_R_{t:.3f}_s", -# save_fig=save_fig, -# ) -# -# if plot_rho: -# x_data = ts_rho -# x_fit = x_data.interp(R=ts_fit.R) -# plot_gpr_fit( -# ts_te, -# ts_fit, -# ts_fit_err, -# t_ts, -# x_data=x_data, -# x_fit=x_fit, -# ylabel="[eV]", -# xlabel="rho-poloidal", -# title=str(st40.pulse), -# label=f"El. (TS) {t_ts:.3f} s", -# color=COLORS["electron"], -# ) -# -# if plot_rho and np.abs(t_ts - t_pi) < 0.02: -# plt.errorbar( -# pi_rho.sel(t=t_pi).values, -# pi_ti.sel(t=t_pi).values, -# pi_ti_err.sel(t=t_pi).values, -# color=COLORS["ion"], -# ) -# plt.plot( -# pi_rho.sel(t=t_pi).values, -# pi_ti.sel(t=t_pi).values, -# marker="s", -# color=COLORS["ion"], -# label=f"Ion (PI) {t_pi:.3f} s", -# ) -# -# if plot_rho and np.abs(t_ts - t_tws) < 0.02: -# plt.errorbar( -# tws_rho.sel(t=t_tws).values, -# tws_ti.sel(t=t_tws).values, -# tws_ti_err.sel(t=t_tws).values, -# color=COLORS["ion"], -# ) -# plt.plot( -# tws_rho.sel(t=t_tws).values, -# tws_ti.sel(t=t_tws).values, -# marker="^", -# color=COLORS["ion"], -# label=f"Ion (TWS) {t_tws:.3f} s", -# ) -# if plot_rho: -# plt.legend(loc="upper left") -# if save_fig: -# save_figure( -# FIG_PATH, -# f"{fig_name}_vs_rho_{t:.3f}_s", -# save_fig=save_fig, -# ) -# -# R0 = 0.48 -# if plot_spectra and np.abs(t_ts - t_pi) < 0.02: -# plt.figure() -# R = pi_spectra.R.sel(R=R0, method="nearest").values -# pi_spectra.sel(t=t_pi, R=R).plot(label="Spectra", color="black") -# pi_fit.sel(t=t_pi, R=R).plot(label="Fit", color="orange") -# plt.legend(loc="upper left") -# plt.title(f"PI spectra {t_tws:.3f} s, R={R:.2f}m") -# -# if plot_spectra and np.abs(t_ts - t_tws) < 0.02: -# plt.figure() -# R = tws_spectra.R.sel(R=R0, method="nearest").values -# tws_spectra.sel(t=t_tws, R=R).plot(label="Spectra", color="black") -# tws_fit.sel(t=t_tws, R=R).plot(label="Fit", color="orange") -# plt.legend(loc="upper left") -# plt.title(f"TWS spectra {t_tws:.3f} s, R={R:.2f}m") -# -# if save_fig: -# plt.close("all") -# else: -# plt.show() -# -# -# if __name__ == "__main__": -# plot_ts_cxrs_profiles() From 62fac27237a5af5d8aec166264564788e36c256f Mon Sep 17 00:00:00 2001 From: Marco Sertoli Date: Fri, 27 Oct 2023 15:23:02 +0100 Subject: [PATCH 04/17] Clean-up of st40 branch approaching move to indica-mcf organisation --- indica/converters/__init__.py | 6 - indica/converters/flux_major_radius.py | 161 --- indica/converters/flux_surfaces.py | 86 -- indica/converters/impact_parameter.py | 286 ---- indica/operators/__init__.py | 10 - indica/operators/bolometry_derivation.py | 538 -------- .../operators/extrapolate_impurity_density.py | 1214 ----------------- indica/operators/impurity_concentration.py | 245 ---- indica/operators/spline_fit.py | 322 ----- indica/operators/spline_fit_easy.py | 4 +- indica/readers/manage_data.py | 204 --- snippets/test_spline_fit.py | 46 - tests/unit/converters/test_flux_surfaces.py | 97 -- tests/unit/data_strategies.py | 200 --- 14 files changed, 2 insertions(+), 3417 deletions(-) delete mode 100644 indica/converters/flux_major_radius.py delete mode 100644 indica/converters/flux_surfaces.py delete mode 100644 indica/converters/impact_parameter.py delete mode 100644 indica/operators/bolometry_derivation.py delete mode 100644 indica/operators/extrapolate_impurity_density.py delete mode 100644 indica/operators/impurity_concentration.py delete mode 100644 indica/operators/spline_fit.py delete mode 100644 indica/readers/manage_data.py delete mode 100644 snippets/test_spline_fit.py delete mode 100644 tests/unit/converters/test_flux_surfaces.py diff --git a/indica/converters/__init__.py b/indica/converters/__init__.py index b0ce5934..b5a6cae6 100644 --- a/indica/converters/__init__.py +++ b/indica/converters/__init__.py @@ -3,9 +3,6 @@ from .abstractconverter import CoordinateTransform from .abstractconverter import EquilibriumException -from .flux_major_radius import FluxMajorRadCoordinates -from .flux_surfaces import FluxSurfaceCoordinates -from .impact_parameter import ImpactParameterCoordinates from .line_of_sight import LineOfSightTransform from .lines_of_sight import LinesOfSightTransform from .time import bin_to_time_labels @@ -16,9 +13,6 @@ __all__ = [ "CoordinateTransform", "EquilibriumException", - "FluxMajorRadCoordinates", - "FluxSurfaceCoordinates", - "ImpactParameterCoordinates", "LinesOfSightTransform", "LineOfSightTransform", "TransectCoordinates", diff --git a/indica/converters/flux_major_radius.py b/indica/converters/flux_major_radius.py deleted file mode 100644 index b2daaa10..00000000 --- a/indica/converters/flux_major_radius.py +++ /dev/null @@ -1,161 +0,0 @@ -"""Defines a coordinate system for use when estimating emissivity data.""" - -from typing import Callable -from typing import cast -from typing import Dict -from typing import Optional - -import numpy as np -from xarray import DataArray - -from .abstractconverter import Coordinates -from .abstractconverter import CoordinateTransform -from .flux_surfaces import FluxSurfaceCoordinates -from ..numpy_typing import LabeledArray - - -class FluxMajorRadCoordinates(CoordinateTransform): - """A coordinate system that uses a flux surface :math:`\\rho` and - major radius to determine spatial positions. This is used, e.g., - when estimating an emissivity profile of the plasma based on X-ray - data. Note that this coordinate system loses information about - vertical position. When converting to (R,z) coordinates, the - results will have ``z >= 0``. - - Parameters - ---------- - flux_surfaces : FluxSurfaceCoordinates - The flux surface coordinate system to use for :math:`\\rho`. - - """ - - _INVERSE_CONVERSION_METHODS: Dict[str, str] = { - "FluxSurfaceCoordinates": "_convert_from_flux_coords" - } - - x2_name = "R" - - def __init__(self, flux_surfaces: FluxSurfaceCoordinates): - self.flux_surfaces = flux_surfaces - self.equilibrium = flux_surfaces.equilibrium - self.flux_kind = flux_surfaces.flux_kind - self.x1_name = flux_surfaces.x1_name - - def get_converter( - self, other: CoordinateTransform, reverse=False - ) -> Optional[Callable[[LabeledArray, LabeledArray, LabeledArray], Coordinates]]: - """Checks if there is a shortcut to convert between these coordiantes, - returning it if so. This can sometimes save the step of - converting to (R, z) coordinates first. - - Parameters - ---------- - other - The other transform whose coordinate system you want to convert to. - reverse - If True, try to return a function which converts _from_ ``other`` - to this coordinate system. - - Returns - ------- - : - If a shortcut function is available, return it. Otherwise, None. - - Note - ---- - Implementations should call ``other.get_converter(self, reverse=True``. For - obvious reasons, however, they should **only do this when - ``reverse == False``**. - - """ - if reverse: - if other == self.flux_surfaces: - return self._convert_from_flux_coords - else: - return None - return other.get_converter(self, True) - - def _convert_from_flux_coords( - self, rho: LabeledArray, theta: LabeledArray, t: LabeledArray - ) -> Coordinates: - """Convert from to a flux coordinate system to this one. - - Parameters - ---------- - rho - The flux surface value. - theta - The poloidal angle on the flux surface. - t - The time coordinate - - Returns - ------- - x1 - The first spatial coordinate in this system. - x2 - The second spatial coordinate in this system. - - """ - R, z = self.flux_surfaces.convert_to_Rz(rho, theta, t) - return rho, R - - def convert_to_Rz( - self, x1: LabeledArray, x2: LabeledArray, t: LabeledArray - ) -> Coordinates: - """Convert from this coordinate to the R-z coordinate system. - - Parameters - ---------- - x1 - The first spatial coordinate in this system. - x2 - The second spatial coordinate in this system. - t - The time coordinate (if there is one, otherwise ``None``) - - Returns - ------- - R - Major radius coordinate - z - Height coordinate - - """ - theta_grid = DataArray( - np.linspace(0.0, np.pi, 100), dims=("theta",), name="theta" - ) - R, z = self.flux_surfaces.convert_to_Rz(x1, theta_grid, t) - theta_vals = cast(DataArray, R).indica.invert_interp(x2, target="theta") - return x2, cast(DataArray, z).indica.interp2d(theta=theta_vals) - - def convert_from_Rz( - self, R: LabeledArray, z: LabeledArray, t: LabeledArray - ) -> Coordinates: - """Convert from the master coordinate system to this coordinate. - - Parameters - ---------- - R - Major radius coordinate - z - Height coordinate - t - Time coordinate) - - Returns - ------- - x1 - The first spatial coordinate in this system. - x2 - The second spatial coordinate in this system. - - """ - rho, theta = self.flux_surfaces.convert_from_Rz(R, z, t) - return rho, R - - def __eq__(self, other: object) -> bool: - if not isinstance(other, self.__class__): - return False - result = self._abstract_equals(other) - return result and self.flux_surfaces == other.flux_surfaces diff --git a/indica/converters/flux_surfaces.py b/indica/converters/flux_surfaces.py deleted file mode 100644 index 7e05f00e..00000000 --- a/indica/converters/flux_surfaces.py +++ /dev/null @@ -1,86 +0,0 @@ -"""Class to handle conversions to and from flux surface coordinates.""" - -from typing import cast -from typing import Optional - -import numpy as np - -from indica.numpy_typing import Coordinates -from .abstractconverter import CoordinateTransform -from ..numpy_typing import LabeledArray - - -class FluxSurfaceCoordinates(CoordinateTransform): - """Class for polar coordinate systems using flux surfaces - for the radial coordinate. - - Parameters - ---------- - kind - The type of flux surface to use. Must be a valid argument for methods - on the :py:class:`indica.equilibrium.Equilibrium` class. - - """ - - x2_name = "theta" - - def __init__( - self, - kind: str, - ): - self.flux_kind = kind - self.x1_name = "rho_" + kind - - def convert_to_Rz( - self, x1: LabeledArray, x2: LabeledArray, t: Optional[LabeledArray] = None - ) -> Coordinates: - """Convert from this coordinate to the R-z coordinate system. - - Parameters - ---------- - x1 - The first spatial coordinate in this system. - x2 - The second spatial coordinate in this system. - t - The time coordinate (if there is one, otherwise ``None``) - - Returns - ------- - R - Major radius coordinate - z - Height coordinate - - """ - return self.equilibrium.spatial_coords(x1, x2, t, self.flux_kind)[0:2] - - def convert_from_Rz( - self, R: LabeledArray, z: LabeledArray, t: Optional[LabeledArray] = None - ) -> Coordinates: - """Convert from the master coordinate system to this coordinate. - - Parameters - ---------- - R - Major radius coordinate - z - Height coordinate - t - Time coordinate) - - Returns - ------- - x1 - The first spatial coordinate in this system. - x2 - The second spatial coordinate in this system. - - """ - return self.equilibrium.flux_coords(R, z, t, self.flux_kind)[0:2] - - def __eq__(self, other: object) -> bool: - if not isinstance(other, self.__class__): - return False - result = self._abstract_equals(other) - return cast(bool, result and np.all(self.flux_kind == other.flux_kind)) diff --git a/indica/converters/impact_parameter.py b/indica/converters/impact_parameter.py deleted file mode 100644 index 7b2aac7c..00000000 --- a/indica/converters/impact_parameter.py +++ /dev/null @@ -1,286 +0,0 @@ -"""Coordinate systems based on volume enclosed by flux surfaces.""" - -from typing import Callable -from typing import cast -from typing import Optional -from typing import Tuple - -import numpy as np -from xarray import DataArray -from xarray import where - -from .abstractconverter import Coordinates -from .abstractconverter import CoordinateTransform -from .flux_surfaces import FluxSurfaceCoordinates -from .lines_of_sight import LinesOfSightTransform -from ..numpy_typing import LabeledArray -from ..utilities import coord_array - - -class ImpactParameterCoordinates(CoordinateTransform): - """Class for coordinate systems based on lines-of-sight, but using the - impact parameter (smallest flux value along a line-of-sight) as the - label for the first coordinate. - - Parameters - ---------- - lines_of_sight - A line-of-sight coordinate system for which this transform will get - the impact parameters. - flux_surfaces - The flux surface coordinate system from which flux values will be used - for the impact parameters. - num_intervals - The number of points along the line of sight at which to evaulate the - flux surface value. - times - The times at which to evaluate the impact parameter. Defaults to all - times at which equilibrium data is available. - """ - - def __init__( - self, - lines_of_sight: LinesOfSightTransform, - flux_surfaces: FluxSurfaceCoordinates, - num_intervals: int = 100, - times: Optional[LabeledArray] = None, - ): - # TODO: Set up proper defaults - self.lines_of_sight = lines_of_sight - self.flux_surfaces = flux_surfaces - if not hasattr(flux_surfaces, "equilibrium"): - raise ValueError( - "Flux surface coordinate system must have an equilibrium set." - ) - self.equilibrium = flux_surfaces.equilibrium - if ( - hasattr(lines_of_sight, "equilibrium") - and lines_of_sight.equilibrium is not flux_surfaces.equilibrium - ): - raise ValueError( - "Two coordinate systems must have the same equilibrium object." - ) - rmag = self.equilibrium.rmag - zmag = self.equilibrium.zmag - self.x1_name = lines_of_sight.x1_name[:-6] + flux_surfaces.x1_name - self.x2_name = lines_of_sight.x2_name - R, z = cast( - Tuple[DataArray, DataArray], - lines_of_sight.convert_to_Rz( - coord_array( - np.arange(len(lines_of_sight.x_start)), lines_of_sight.x1_name - ), - coord_array(np.linspace(0.0, 1.0, num_intervals + 1), self.x2_name), - 0.0, - ), - ) - rho, _ = cast( - Tuple[DataArray, DataArray], flux_surfaces.convert_from_Rz(R, z, times) - ) - rho = where(rho < 0, float("nan"), rho) - t = rho.coords["t"] - loc = rho.argmin(self.x2_name) - theta = np.arctan2( - z.sel({self.x2_name: 0.0}).mean() - - np.mean(lines_of_sight._machine_dims[1]), - R.sel({self.x2_name: 0.0}).mean() - - np.mean(lines_of_sight._machine_dims[0]), - ) - if np.pi / 4 <= np.abs(theta) <= 3 * np.pi / 4: - sign = where( - R.isel({self.x2_name: loc}) < rmag.interp(t=t, method="nearest"), -1, 1 - ) - else: - sign = where( - z.isel({self.x2_name: loc}) < zmag.interp(t=t, method="nearest"), -1, 1 - ) - self.rho_min = sign * rho.isel({self.x2_name: loc}) - - def get_converter( - self, other: "CoordinateTransform", reverse=False - ) -> Optional[Callable[[LabeledArray, LabeledArray, LabeledArray], Coordinates]]: - """Checks if there is a shortcut to convert between these coordiantes, - returning it if so. This can sometimes save the step of - converting to (R, z) coordinates first. - - Parameters - ---------- - other - The other transform whose coordinate system you want to convert to. - reverse - If True, try to return a function which converts _from_ ``other`` - to this coordinate system. - - Returns - ------- - : - If a shortcut function is available, return it. Otherwise, None. - - Note - ---- - Implementations should call ``other.get_converter(self, reverse=True``. For - obvious reasons, however, they should **only do this when - ``reverse == False``**. - - """ - if reverse: - if other == self.lines_of_sight: - return self._convert_from_los - else: - return None - if other == self.lines_of_sight: - return self._convert_to_los - else: - return other.get_converter(self, True) - - def _convert_to_los( - self, min_rho: LabeledArray, x2: LabeledArray, t: LabeledArray - ) -> Coordinates: - """Convert from this coordinate system to a line-of-sight coordinate system. - - Parameters - ---------- - min_rho - The first spatial coordinate in this system. - x2 - The second spatial coordinate in this system. - t - The time coordinate - - Returns - ------- - los - Index of the line of sight which that impact parameter corresponds - to. - theta - Position along the line of sight. - - """ - # Cubic splines aren't guaranteed to be monotonicuse linear. - # TODO: Find a better spline that I can ensure is monotonic - return ( - self.rho_min.interp(t=t, method="nearest").indica.invert_interp( - min_rho, self.lines_of_sight.x1_name, method="linear" - ), - x2, - ) - - def _convert_from_los( - self, x1: LabeledArray, x2: LabeledArray, t: LabeledArray - ) -> Coordinates: - """Converts from line of sight coordinates to this coordinate system. - - Parameters - ---------- - x1 - The index for the line of sight. - x2 - Position along the line of sight. - t - The time coordinate - - Returns - ------- - min_rho - Lowest flux surface value the line of sight touches. - x2 - Position along the line of sight. - - """ - # Cubic splines aren't guaranteed to be monotonic, so use linear - # TODO: Find a better spline that I can ensure is monotonic - return ( - self.rho_min.interp(t=t, method="nearest").indica.interp2d( - {self.lines_of_sight.x1_name: x1}, method="linear" - ), - x2, - ) - - def convert_to_Rz( - self, x1: LabeledArray, x2: LabeledArray, t: LabeledArray - ) -> Coordinates: - """Convert from this coordinate to the R-z coordinate system. - - Parameters - ---------- - x1 - The first spatial coordinate in this system. - x2 - The second spatial coordinate in this system. - t - The time coordinate (if there is one, otherwise ``None``) - - Returns - ------- - R - Major radius coordinate - z - Height coordinate - - """ - index, position = self._convert_to_los(x1, x2, t) - return self.lines_of_sight.convert_to_Rz(index, position, t) - - def convert_from_Rz( - self, R: LabeledArray, z: LabeledArray, t: LabeledArray - ) -> Coordinates: - """Convert from the master coordinate system to this coordinate. - - Parameters - ---------- - R - Major radius coordinate - z - Height coordinate - t - Time coordinate - - Returns - ------- - x1 - The first spatial coordinate in this system. - x2 - The second spatial coordinate in this system. - - """ - x1, x2 = self.lines_of_sight.convert_from_Rz(R, z, t) - return self._convert_from_los(x1, x2, t) - - def distance( - self, - direction: str, - x1: LabeledArray, - x2: LabeledArray, - t: LabeledArray, - ) -> LabeledArray: - """Implementation of calculation of physical distances between points - in this coordinate system. This accounts for potential toroidal skew of - lines. - - """ - return self.lines_of_sight.distance( - direction, *self._convert_to_los(x1, x2, t), t - ) - - def drho(self) -> float: - """Calculates the average difference in impact parameters between - adjacent lines of sight.""" - drhos = np.abs( - self.rho_min.isel({self.x1_name: slice(1, None)}).data - - self.rho_min.isel({self.x1_name: slice(None, -1)}).data - ) - return np.mean(drhos) - - def rhomax(self) -> float: - """Returns the time-averaged maximum impact parameter on the low flux - surface. - - """ - return self.rho_min.mean("t").max() - - def __eq__(self, other: object) -> bool: - if not isinstance(other, self.__class__): - return False - result = self._abstract_equals(other) - result = result and self.flux_surfaces == other.flux_surfaces - return result and self.lines_of_sight == other.lines_of_sight diff --git a/indica/operators/__init__.py b/indica/operators/__init__.py index c79eca90..4bd71ea9 100644 --- a/indica/operators/__init__.py +++ b/indica/operators/__init__.py @@ -2,14 +2,9 @@ from .abstractoperator import OperatorError from .atomic_data import FractionalAbundance from .atomic_data import PowerLoss -from .bolometry_derivation import BolometryDerivation from .centrifugal_asymmetry import AsymmetryParameter from .centrifugal_asymmetry import ToroidalRotation -from .extrapolate_impurity_density import ExtrapolateImpurityDensity -from .impurity_concentration import ImpurityConcentration from .mean_charge import MeanCharge -from .spline_fit import Spline -from .spline_fit import SplineFit from .zeff import CalcZeff __all__ = [ @@ -17,13 +12,8 @@ "OperatorError", "FractionalAbundance", "PowerLoss", - "BolometryDerivation", "ToroidalRotation", "AsymmetryParameter", - "ExtrapolateImpurityDensity", - "ImpurityConcentration", "MeanCharge", - "Spline", - "SplineFit", "CalcZeff", ] diff --git a/indica/operators/bolometry_derivation.py b/indica/operators/bolometry_derivation.py deleted file mode 100644 index a0a98aa9..00000000 --- a/indica/operators/bolometry_derivation.py +++ /dev/null @@ -1,538 +0,0 @@ -from typing import cast -from typing import List -from typing import Sequence -from typing import Tuple -from typing import Union - -import numpy as np -from xarray import concat -from xarray import DataArray -from xarray.core.common import zeros_like - -from indica.converters.flux_surfaces import FluxSurfaceCoordinates -from indica.converters.lines_of_sight import LinesOfSightTransform -from indica.operators.main_ion_density import MainIonDensity -from indica.operators.mean_charge import MeanCharge -from .abstractoperator import EllipsisType -from .abstractoperator import Operator -from .. import session -from ..datatypes import DataType -from ..utilities import input_check - - -class BolometryDerivation(Operator): - """Class to hold relevant variables and functions relating to the derivation of - bolometry data from plasma quantities (densities, temperatures, etc.) - - Parameters - ---------- - flux_surfaces - FluxSurfaceCoordinates object representing polar coordinate systems - using flux surfaces for the radial coordinate. - LoS_bolometry_data - Line-of-sight bolometry data in the same format as given in: - tests/unit/operator/KB5_Bolometry_data.py - t_arr - Array of time values to interpolate the (rho, theta) grids on. - xarray.DataArray with dimensions (t). - impurity_densities - Densities for all impurities - (including the extrapolated smooth density of the impurity in question), - xarray.DataArray with dimensions (elements, rho, theta, t). - frac_abunds - Fractional abundances list of fractional abundances - (an xarray.DataArray for each impurity) dimensions of each element in - the list are (ion_charges, rho, t). - impurity_elements - List of element symbols(as strings) for all impurities. - electron_density - xarray.DataArray of electron density, xarray.DataArray wit dimensions (rho, t) - main_ion_power_loss - Power loss associated with the main ion (eg. deuterium), - xarray.DataArray with dimensions (rho, t) - main_ion_density - Density profile for the main ion, - xarray.DataArray with dimensions (rho, theta, t) - - Returns - ------- - derived_power_loss_LoS_tot - Total derived bolometric power loss values along all lines-of-sight. - xarray.DataArray with dimensions (channels, t) or (channels) depending - on whether t_val is provided. - - Attributes - ---------- - ARGUMENT_TYPES: List[DataType] - Ordered list of the types of data expected for each argument of the - operator. - RESULT_TYPES: List[DataType] - Ordered list of the types of data returned by the operator. - """ - - ARGUMENT_TYPES: List[Union[DataType, EllipsisType]] = [ - ("plasma", "flux_surface_coordinates"), - ("bolometric", "lines_of_sight_data"), - ("plasma", "times"), - ("impurities", "number_density"), - ("impurities", "fractional_abundance"), - ("impurities", "elements"), - ("electrons", "number_density"), - ("main_ion", "total radiated power loss"), - ("main_ion", "number_density"), - ] - - RESULT_TYPES: List[Union[DataType, EllipsisType]] = [] - - def __init__( - self, - flux_surfs: FluxSurfaceCoordinates, - LoS_bolometry_data: Sequence, - t_arr: DataArray, - impurity_densities: DataArray, - frac_abunds: Sequence, - impurity_elements: Sequence, - electron_density: DataArray, - main_ion_power_loss: DataArray, - impurities_power_loss: DataArray, - sess: session.Session = session.global_session, - ): - """Initialises the BolometryDerivation class. Checks all inputs for errors.""" - super().__init__(sess=sess) - - input_check( - "flux_surfs", - flux_surfs, - FluxSurfaceCoordinates, - ) - - input_check("LoS_bolometry_data", LoS_bolometry_data, Sequence) - - input_check("t_arr", t_arr, DataArray, ndim_to_check=1, strictly_positive=False) - - input_check( - "impurity_densities", - impurity_densities, - DataArray, - ndim_to_check=4, - strictly_positive=False, - ) - - input_check("frac_abunds", frac_abunds, Sequence) - - input_check("impurity_elements", impurity_elements, Sequence) - - input_check( - "electron_density", - electron_density, - DataArray, - ndim_to_check=2, - strictly_positive=False, - ) - - input_check( - "main_ion_power_loss", - main_ion_power_loss, - DataArray, - ndim_to_check=2, - strictly_positive=False, - ) - - input_check( - "impurities_power_loss", - impurities_power_loss, - DataArray, - ndim_to_check=3, - strictly_positive=False, - ) - - self.flux_surfaces = flux_surfs - self.LoS_bolometry_data = LoS_bolometry_data - self.t_arr = t_arr - self.impurity_densities = impurity_densities - self.frac_abunds = frac_abunds - self.impurity_elements = impurity_elements - self.electron_density = electron_density - self.main_ion_power_loss = main_ion_power_loss - self.impurities_power_loss = impurities_power_loss - - def return_types(self, *args: DataType) -> Tuple[DataType, ...]: - return super().return_types(*args) - - def __bolometry_coord_transforms(self): - """Transform the bolometry coords from LoS to (rho, theta) and (R, z). - - Returns - ------- - LoS_coords - List of dictionaries containing the rho, theta, x and z arrays - and dl for the resolution of the LoS coordinates. - """ - LoS_coords = [] - for iLoS in range(len(self.LoS_bolometry_data)): - LoS_transform = LinesOfSightTransform(*self.LoS_bolometry_data[iLoS]) - - x1_name = LoS_transform.x1_name - x2_name = LoS_transform.x2_name - - x1 = DataArray( - data=np.array([0]), coords={x1_name: np.array([0])}, dims=[x1_name] - ) - x2 = DataArray( - data=np.linspace(0, 1, 30), - coords={x2_name: np.linspace(0, 1, 30)}, - dims=[x2_name], - ) - - R_arr, z_arr = LoS_transform.convert_to_Rz(x1, x2, self.t_arr) - - rho_arr, theta_arr = self.flux_surfaces.convert_from_Rz(R_arr, z_arr) - rho_arr = cast(DataArray, rho_arr).interp(t=self.t_arr, method="linear") - theta_arr = cast(DataArray, theta_arr).interp(t=self.t_arr, method="linear") - - rho_arr = np.abs(rho_arr) - rho_arr = rho_arr.assign_coords({x2_name: x2}) - rho_arr = rho_arr.drop(x1_name).squeeze() - rho_arr = rho_arr.fillna(2.0) - theta_arr = theta_arr.drop(x1_name).squeeze() - - dl = LoS_transform.distance(x2_name, DataArray(0), x2[0:2], 0) - - LoS_coords.append( - dict( - { - # dimensions for rho_arr and - # theta_arr are (channel no., distance along LoS) - "rho_poloidal": rho_arr, - "theta": theta_arr, - # dimensions for dl are (channel no.) - "dl": dl, - # dimensions for R_arr and - # z_arr are (channel no., distance along LoS) - "R": R_arr, - "z": z_arr, - # dimensions for t are (t) - "t": self.t_arr, - "label": self.LoS_bolometry_data[iLoS][6], - } - ) - ) - - self.LoS_coords = LoS_coords - - return LoS_coords - - def __bolometry_setup(self): - """Calculating main ion density for the bolometry derivation. - - Returns - ------- - main_ion_density - Density profile for the main ion, - xarray.DataArray with dimensions (rho, theta, t) - """ - mean_charges = zeros_like(self.electron_density) - mean_charges = mean_charges.data - mean_charges = np.tile(mean_charges, (len(self.impurity_elements), 1, 1)) - # Ignoring mypy error since mypy refuses to acknowlege electron_density.coords - # as a dictionary - mean_charges_coords = { - "element": self.impurity_elements, # type: ignore - **self.electron_density.coords, # type: ignore - } - mean_charges = DataArray( - data=mean_charges, - coords=mean_charges_coords, # type:ignore - dims=["element", *self.electron_density.dims], - ) - - for ielement, element in enumerate(self.impurity_elements): - mean_charge = MeanCharge() - mean_charge = mean_charge(self.frac_abunds[ielement], element) - mean_charges.loc[element] = mean_charge - - main_ion_density_obj = MainIonDensity() - main_ion_density = main_ion_density_obj( - self.impurity_densities, self.electron_density, mean_charges - ) - - main_ion_density = main_ion_density.transpose("rho_poloidal", "theta", "t") - - self.main_ion_density = main_ion_density - - return main_ion_density - - def __bolometry_channel_filter(self): - """Filters the bolometry data to reduce the number of channels by eliminating - channels that are too close together. - - Returns - ------- - LoS_bolometry_data_trimmed - Lines-of-sight data with a reduced number of channels. List with - the same formatting as self.LoS_bolometry_data. - LoS_coords_trimmed - LoS coordinates (as returned by bolometry_coord_transforms) with a - reduced number of channels. List with the same formatting as - self.LoS_coords. - """ - LoS_bolometry_data_in = self.LoS_bolometry_data - LoS_coords_in = self.LoS_coords - - LoS_bolometry_data = [] - LoS_coords = [] - - t_arr = LoS_coords_in[0]["t"] - - R_central = self.flux_surfaces.equilibrium.rmag.interp(t=t_arr, method="linear") - - z_central = self.flux_surfaces.equilibrium.zmag.interp(t=t_arr, method="linear") - - LoS_R_midplane_multi = [] - - for iLoS in range(len(LoS_bolometry_data_in)): - R_arr = LoS_coords_in[iLoS]["R"] - z_arr = LoS_coords_in[iLoS]["z"] - - try: - midplane_LoS_pos = z_arr.indica.invert_interp( - values=z_central.values[0], target=z_arr.dims[1], method="nearest" - ) - except ValueError: - midplane_LoS_pos = 2.0 - - LoS_R_midplane = R_arr.interp( - {R_arr.dims[1]: midplane_LoS_pos}, method="nearest" - ) - - if (LoS_R_midplane > R_central).all(): - LoS_bolometry_data.append(LoS_bolometry_data_in[iLoS]) - LoS_coords.append(LoS_coords_in[iLoS]) - LoS_R_midplane_multi.append(LoS_R_midplane.data[0]) - - LoS_R_midplane_positions = DataArray( - data=LoS_R_midplane_multi, - coords={"channels": [j[6] for j in LoS_bolometry_data]}, - dims=["channels"], - ) - - LoS_R_midplane_positions_sorted = LoS_R_midplane_positions.sortby( - LoS_R_midplane_positions - ) - - radial_resolution_threshold = 0.1 # in metres - - LoS_R_midplane_positions_diff = LoS_R_midplane_positions_sorted.diff( - "channels", label="upper" - ) - - LoS_R_midplane_positions_diff_first = DataArray( - data=np.array([True]), - coords={ - "channels": np.array( - [LoS_R_midplane_positions_sorted.coords["channels"].data[0]] - ) - }, - dims=["channels"], - ) - - LoS_R_midplane_positions_diff = concat( - [LoS_R_midplane_positions_diff_first, LoS_R_midplane_positions_diff], - dim="channels", - ) - - LoS_R_midplane_trim_mask = ( - LoS_R_midplane_positions_diff > radial_resolution_threshold - ) - - LoS_R_midplane_trimmed = LoS_R_midplane_positions_sorted[ - LoS_R_midplane_trim_mask - ] - - LoS_bolometry_data_trimmed = [] - LoS_coords_trimmed = [] - for iLoS in range(len(LoS_bolometry_data)): - orig_channel = LoS_bolometry_data[iLoS][6] - if orig_channel in LoS_R_midplane_trimmed.coords["channels"].data: - LoS_bolometry_data_trimmed.append(LoS_bolometry_data[iLoS]) - LoS_coords_trimmed.append(LoS_coords[iLoS]) - - self.LoS_bolometry_data_trimmed, self.LoS_coords_trimmed = ( - LoS_bolometry_data_trimmed, - LoS_coords_trimmed, - ) - - return LoS_bolometry_data_trimmed, LoS_coords_trimmed - - def __bolometry_derivation( - self, - trim: bool = False, - t_val: float = None, - ): - """Derive bolometry including the extrapolated smoothed impurity density. - - Parameters - ---------- - trim - Boolean specifying whether the number of bolometry channels are trimmed. - t_val - Optional time value(float) for which to calculate the bolometry data. - - Returns - ------- - derived_power_loss_LoS_tot - Total derived bolometric power loss values along all lines-of-sight. - xarray.DataArray with dimensions (channels, t) or (channels) depending - on whether t_val is provided. - """ - if trim: - if not ( - hasattr(self, "LoS_bolometry_data_trimmed") - and hasattr(self, "LoS_coords_trimmed") - ): - raise AttributeError( - 'Argument "trim" is set to True but bolometry_channel_filter() \ - has not yet been run at least once.' - ) - LoS_bolometry_data = self.LoS_bolometry_data_trimmed - LoS_coords_in = self.LoS_coords_trimmed - else: - LoS_bolometry_data = self.LoS_bolometry_data - LoS_coords_in = self.LoS_coords - - if t_val is not None: - LoS_coords = cast(Sequence, [{} for i in LoS_coords_in]) - impurity_densities = self.impurity_densities.sel(t=t_val) - main_ion_power_loss = self.main_ion_power_loss.sel(t=t_val) - impurities_power_loss = self.impurities_power_loss.sel(t=t_val) - electron_density = self.electron_density.sel(t=t_val) - main_ion_density = self.main_ion_density.sel(t=t_val) - for icoord in range(len(LoS_coords_in)): - LoS_coords[icoord]["rho_poloidal"] = LoS_coords_in[icoord][ - "rho_poloidal" - ].sel(t=t_val) - LoS_coords[icoord]["theta"] = LoS_coords_in[icoord]["theta"].sel( - t=t_val - ) - LoS_coords[icoord]["dl"] = LoS_coords_in[icoord]["dl"] - LoS_coords[icoord]["R"] = LoS_coords_in[icoord]["R"] - LoS_coords[icoord]["z"] = LoS_coords_in[icoord]["z"] - else: - LoS_coords = LoS_coords_in - impurity_densities = self.impurity_densities - main_ion_power_loss = self.main_ion_power_loss - impurities_power_loss = self.impurities_power_loss - electron_density = self.electron_density - main_ion_density = self.main_ion_density - - derived_power_loss = electron_density * (main_ion_density * main_ion_power_loss) - impurities_losses = impurity_densities * impurities_power_loss - impurities_losses = impurities_losses.sum(dim="element") - impurities_losses *= electron_density - derived_power_loss += impurities_losses - - if t_val is not None: - derived_power_loss = derived_power_loss.transpose("rho_poloidal", "theta") - - derived_power_loss_LoS_tot = DataArray( - data=np.zeros((len(LoS_bolometry_data))), - coords={ - "channels": np.linspace( - 0, - len(LoS_bolometry_data), - len(LoS_bolometry_data), - endpoint=False, - ), - }, - dims=["channels"], - ) - - else: - derived_power_loss = derived_power_loss.transpose( - "t", "rho_poloidal", "theta" - ) - - t_arr = derived_power_loss.coords["t"] - - derived_power_loss_LoS_tot = DataArray( - data=np.zeros((len(LoS_bolometry_data), t_arr.shape[0])), - coords={ - "channels": np.linspace( - 0, - len(LoS_bolometry_data), - len(LoS_bolometry_data), - endpoint=False, - ), - "t": t_arr, - }, - dims=["channels", "t"], - ) - - for iLoS in range(len(LoS_bolometry_data)): - LoS_transform = LinesOfSightTransform(*LoS_bolometry_data[iLoS]) - - x2_name = LoS_transform.x2_name - - rho_arr = LoS_coords[iLoS]["rho_poloidal"] - theta_arr = LoS_coords[iLoS]["theta"] - - derived_power_loss_LoS = derived_power_loss.interp( - {"rho_poloidal": rho_arr, "theta": theta_arr} - ) - - derived_power_loss_LoS = derived_power_loss_LoS.fillna(0.0) - - dl = LoS_coords[iLoS]["dl"] - dl = cast(DataArray, dl)[1] - - derived_power_loss_LoS = derived_power_loss_LoS.sum(dim=x2_name) * dl - derived_power_loss_LoS_tot[iLoS] = derived_power_loss_LoS.squeeze() - - derived_power_loss_LoS_tot = derived_power_loss_LoS_tot.fillna(0.0) - - self.derived_power_loss_LoS_tot = derived_power_loss_LoS_tot - - return derived_power_loss_LoS_tot - - def __call__( # type: ignore - self, deriv_only: bool = False, trim: bool = False, t_val: float = None - ): - """Varying workflow to derive bolometry from plasma quantities. - (Varying as in, if full setup and derivation is needed or only derivaiton.) - - Parameters - ---------- - deriv_only - Optional boolean specifying if only derivation is needed(True) or if full - setup and derivation is needed(False). - trim - Optional boolean specifying whether to use bolometry data with trimmed - channels(True) or not(False). - t_val - Optional time value for which to calculate the bolometry data. - (This is passed to bolometry_derivation()) - - Returns - ------- - derived_power_loss_LoS_tot - Total derived bolometric power loss values along all lines-of-sight. - xarray.DataArray with dimensions (channels, t) or (channels) depending - on whether t_val is provided. - """ - input_check("deriv_only", deriv_only, bool) - input_check("trim", trim, bool) - - if t_val is not None: - input_check("t_val", t_val, float, strictly_positive=False) - - if not deriv_only: - self.__bolometry_coord_transforms() - - self.__bolometry_setup() - - if trim: - self.__bolometry_channel_filter() - - self.__bolometry_derivation(trim, t_val) - - return self.derived_power_loss_LoS_tot diff --git a/indica/operators/extrapolate_impurity_density.py b/indica/operators/extrapolate_impurity_density.py deleted file mode 100644 index 4876db6f..00000000 --- a/indica/operators/extrapolate_impurity_density.py +++ /dev/null @@ -1,1214 +0,0 @@ -from typing import cast -from typing import List -from typing import Sequence -from typing import Tuple -from typing import Union - -import numpy as np -from scipy.interpolate import UnivariateSpline -from scipy.optimize import least_squares -from scipy.stats import norm -from xarray import concat -from xarray import DataArray -from xarray.core.common import zeros_like - -from indica.converters.flux_surfaces import FluxSurfaceCoordinates -from indica.operators.bolometry_derivation import BolometryDerivation -from .abstractoperator import EllipsisType -from .abstractoperator import Operator -from .. import session -from ..datatypes import DataType -from ..utilities import input_check - - -def asymmetry_from_R_z( - data_R_z: DataArray, - flux_surfaces: FluxSurfaceCoordinates, - rho_arr: DataArray, - threshold_rho: DataArray = None, - t_arr: DataArray = None, -): - """Function to calculate an asymmetry parameter from a given density profile in - (R, z, t) coordinates. - - Parameters - ---------- - data_R_z - High-z density profile which is to be used to calculate the asymmetry parameter. - xarray.DataArray with dimensions (R, z, t) - flux_surfaces - FluxSurfaceCoordinates object representing polar coordinate systems - using flux surfaces for the radial coordinate. - rho_arr - 1D xarray.DataArray of rho from 0 to 1. - threshold_rho - rho value denoting the cutoff point beyond which soft x-ray diagnostics - are invalid. It's also used in setting the derived asymmetry parameter to be - flat in the invalid region. - xarray.DataArray with dimensions (t) - t_arr - 1D xarray.DataArray of t. - - Returns - ------- - derived_asymmetry_parameter - Derived asymmetry parameter. xarray.DataArray with dimensions (rho, t) - """ - - input_check("data_R_z", data_R_z, DataArray, 3, strictly_positive=False) - - input_check("flux_surfaces", flux_surfaces, FluxSurfaceCoordinates) - - input_check("rho_arr", rho_arr, DataArray, 1, strictly_positive=False) - - if threshold_rho is not None: - input_check( - "threshold_rho", threshold_rho, DataArray, 1, strictly_positive=True - ) - - if t_arr is None: - t_arr = data_R_z.coords["t"] - else: - input_check("t_arr", t_arr, DataArray, 1, strictly_positive=False) - - theta_arr_ = np.array([0.0, np.pi], dtype=float) - theta_arr = DataArray(data=theta_arr_, coords={"theta": theta_arr_}, dims=["theta"]) - - R_deriv, z_deriv = flux_surfaces.convert_to_Rz(rho_arr, theta_arr) - R_deriv = cast(DataArray, R_deriv).interp(t=t_arr, method="linear") - z_deriv = cast(DataArray, z_deriv).interp(t=t_arr, method="linear") - - R_deriv = cast(DataArray, R_deriv).transpose("rho_poloidal", "theta", "t") - z_deriv = cast(DataArray, z_deriv).transpose("rho_poloidal", "theta", "t") - - data_rho_theta = data_R_z.indica.interp2d( - {"R": R_deriv, "z": z_deriv}, method="linear", assume_sorted=True - ) - data_rho_theta = data_rho_theta.transpose("rho_poloidal", "theta", "t") - - R_lfs_midplane = cast(DataArray, R_deriv).isel(theta=0) # theta = 0.0 - R_hfs_midplane = cast(DataArray, R_deriv).isel(theta=1) # theta = np.pi - - derived_asymmetry_parameter = np.log( - data_rho_theta.isel(theta=1) / data_rho_theta.isel(theta=0) - ) - - derived_asymmetry_parameter /= R_hfs_midplane**2 - R_lfs_midplane**2 - - # Set constant asymmetry parameter for rho<0.1 - derived_asymmetry_parameter = derived_asymmetry_parameter.where( - derived_asymmetry_parameter.coords["rho_poloidal"] > 0.1, - other=derived_asymmetry_parameter.sel({"rho_poloidal": 0.1}, method="nearest"), - ) - - if threshold_rho is not None: - for ind_t, it in enumerate(threshold_rho.coords["t"]): - derived_asymmetry_parameter.loc[ - threshold_rho[ind_t] :, it # type:ignore - ] = derived_asymmetry_parameter.loc[threshold_rho[ind_t], it] - - return derived_asymmetry_parameter - - -def asymmetry_from_rho_theta( - data_rho_theta: DataArray, - flux_surfaces: FluxSurfaceCoordinates, - threshold_rho: DataArray = None, - t_arr: DataArray = None, -): - """Function to calculate an asymmetry parameter from a given density profile in - (rho_poloidal, theta, t) coordinates. - - Parameters - ---------- - data_rho_theta - High-z density profile which is to be used to calculate the asymmetry parameter. - xarray.DataArray with dimensions (rho_poloidal, theta, t) - flux_surfaces - FluxSurfaceCoordinates object representing polar coordinate systems - using flux surfaces for the radial coordinate. - threshold_rho - rho value denoting the cutoff point beyond which soft x-ray diagnostics - are invalid. It's also used in setting the derived asymmetry parameter to be - flat in the invalid region. - xarray.DataArray with dimensions (t) - t_arr - 1D xarray.DataArray of t. - - Returns - ------- - derived_asymmetry_parameter - Derived asymmetry parameter. xarray.DataArray with dimensions (rho, t) - """ - - input_check("data_rho_theta", data_rho_theta, DataArray, 3, strictly_positive=False) - - input_check("flux_surfaces", flux_surfaces, FluxSurfaceCoordinates) - - if threshold_rho is not None: - input_check( - "threshold_rho", threshold_rho, DataArray, 1, strictly_positive=False - ) - - if t_arr is None: - t_arr = data_rho_theta.coords["t"] - else: - input_check("t_arr", t_arr, DataArray, 1, strictly_positive=False) - - rho_arr = data_rho_theta.coords["rho_poloidal"] - theta_arr_ = np.array([0.0, np.pi], dtype=float) - theta_arr = DataArray(data=theta_arr_, coords={"theta": theta_arr_}, dims=["theta"]) - - R_deriv, z_deriv = flux_surfaces.convert_to_Rz(rho_arr, theta_arr) - - R_deriv = cast(DataArray, R_deriv).interp(t=t_arr, method="linear") - z_deriv = cast(DataArray, z_deriv).interp(t=t_arr, method="linear") - - R_deriv = cast(DataArray, R_deriv).transpose("rho_poloidal", "theta", "t") - z_deriv = cast(DataArray, z_deriv).transpose("rho_poloidal", "theta", "t") - - R_lfs_midplane = cast(DataArray, R_deriv).isel(theta=0) # theta = 0.0 - R_hfs_midplane = cast(DataArray, R_deriv).isel(theta=1) # theta = np.pi - - derived_asymmetry_parameter = np.log( - data_rho_theta.interp(theta=np.pi, method="linear") - / data_rho_theta.interp(theta=0.0, method="linear") - ) - - derived_asymmetry_parameter /= R_hfs_midplane**2 - R_lfs_midplane**2 - - # Set constant asymmetry parameter for rho<0.1 - derived_asymmetry_parameter = derived_asymmetry_parameter.where( - derived_asymmetry_parameter.coords["rho_poloidal"] > 0.1, - other=derived_asymmetry_parameter.sel({"rho_poloidal": 0.1}, method="nearest"), - ) - - derived_asymmetry_parameter = derived_asymmetry_parameter.transpose( - "rho_poloidal", "t" - ) - - if threshold_rho is not None: - threshold_asymmetry = derived_asymmetry_parameter.sel( - rho_poloidal=threshold_rho, method="nearest" - ) - derived_asymmetry_parameter = derived_asymmetry_parameter.where( - derived_asymmetry_parameter.rho_poloidal < threshold_rho, - other=threshold_asymmetry, - ) - - return derived_asymmetry_parameter - - -def asymmetry_modifier_from_parameter( - asymmetry_parameter: DataArray, R_deriv: DataArray -): - """ - Convert an asymmetry parameter to an asymmetry modifier. - - Parameters - ---------- - asymmetry_parameter - Asymmetry parameter on (rho, t) - R_deriv - Variable describing value of R for every coordinate on a (rho, theta) grid. - xarray.DataArray with dimensions (rho, theta, t) - - Returns - ------- - asymmetry_modifier - Asymmetry modifier used to transform a low-field side only rho-profile - of a poloidally asymmetric quantity to a full poloidal cross-sectional - profile ie. (rho, t) -> (rho, theta, t). Also can be defined as: - exp(asymmetry_parameter * (R ** 2 - R_lfs ** 2)), where R is the major - radius as a function of (rho, theta, t) and R_lfs is the low-field-side - major radius as a function of (rho, t). xarray DataArray with dimensions - (rho, theta, t) - """ - R_lfs_midplane = cast(DataArray, R_deriv).sel(theta=0, method="nearest") - - asymmetry_modifier = np.exp( - asymmetry_parameter * (R_deriv**2 - R_lfs_midplane**2) - ) - - asymmetry_modifier = asymmetry_modifier.transpose("rho_poloidal", "theta", "t") - - return asymmetry_modifier - - -def recover_threshold_rho(truncation_threshold: float, electron_temperature: DataArray): - """Recover the rho value for a given electron temperature threshold, as in - at what rho location does the electron temperature drop below the specified - temperature. - - Parameters - ---------- - truncation_threshold - User-specified (float) temperature truncation threshold. - electron_temperature - xarray.DataArray of the electron temperature profile, - with dimensions (rho, t). - - Returns - ------- - threshold_rho - rho value beyond which the electron temperature falls below - the truncation_threshold. - """ - input_check( - "truncation_threshold", - truncation_threshold, - float, - strictly_positive=True, - ) - - input_check( - "electron_temperature", - electron_temperature, - DataArray, - 2, - strictly_positive=True, - ) - - try: - assert set(["rho_poloidal"]).issubset( - set(list(electron_temperature.coords.keys())) - ) - except AssertionError: - raise AssertionError("Electron temperature must be a profile of rho.") - - threshold_temp = electron_temperature.where( - (electron_temperature - truncation_threshold >= 0), drop=True - ).min("rho_poloidal") - - threshold_rho_ind = electron_temperature.where( - electron_temperature >= threshold_temp, np.inf - ).argmin("rho_poloidal") - - threshold_rho = electron_temperature.coords["rho_poloidal"].isel( - rho_poloidal=threshold_rho_ind - ) - - return threshold_rho - - -class ExtrapolateImpurityDensity(Operator): - """Extrapolate the impurity density beyond the limits of SXR (Soft X-ray) - - Attributes - ---------- - ARGUMENT_TYPES: List[DataType] - Ordered list of the types of data expected for each argument of the - operator. - RESULT_TYPES: List[DataType] - Ordered list of the types of data returned by the operator. - - Returns - ------- - extrapolated_smooth_density_Rz - Extrapolated and smoothed impurity density, - xarray.DataArray with dimensions (R, z, t). - extrapolated_smooth_density_rho_theta - Extrapolated and smoothed impurity density, - xarray.DataArray with dimensions (rho, theta, t). - t - If ``t`` was not specified as an argument for the __call__ function, - return the time the results are given for. - Otherwise return the argument. - - """ - - ARGUMENT_TYPES: List[Union[DataType, EllipsisType]] = [] - - RESULT_TYPES: List[Union[DataType, EllipsisType]] = [ - ("number_density", "impurity_element"), - ("number_density", "impurity_element"), - ("time", "impurity_element"), - ] - - def __init__( - self, - sess: session.Session = session.global_session, - ): - """Initialises ExtrapolateImpurityDensity class.""" - super().__init__(sess=sess) - - self.halfwidthhalfmax_coeff = np.sqrt(2 * np.log(2)) - self.fullwidthhalfmax_coeff = 2 * self.halfwidthhalfmax_coeff - - def return_types(self, *args: DataType) -> Tuple[DataType, ...]: - return super().return_types(*args) - - def transform_to_rho_theta( - self, - data_R_z: DataArray, - flux_surfaces: FluxSurfaceCoordinates, - rho_arr: DataArray, - t_arr: DataArray = None, - ): - """Function to transform data from an (R, z) grid to a (rho_poloidal, theta) grid - - Parameters - ---------- - data_R_z - xarray.DataArray to be transformed. Dimensions (R, z, t) - flux_surfaces - FluxSurfaceCoordinates object representing polar coordinate systems - using flux surfaces for the radial coordinate. - rho_arr - 1D xarray.DataArray of rho_poloidal from 0 to 1. - t_arr - 1D xarray.DataArray of t. - - Returns - ------- - data_rho_theta - Transformed xarray.DataArray. Dimensions (rho_poloidal, theta, t) - R_deriv - Variable describing value of R in every coordinate on a (rho, theta) grid. - xarray.DataArray with dimensions (rho, theta, t) - z_deriv - Variable describing value of z in every coordinate on a (rho, theta) grid. - xarray.DataArray with dimensions (rho, theta, t) - """ - if t_arr is None: - t_arr = data_R_z.coords["t"] - - theta_arr = np.linspace(-np.pi, np.pi, 21) - # mypy doesn't like re-assignments which changes types. - theta_arr = DataArray( # type: ignore - data=theta_arr, coords={"theta": theta_arr}, dims=["theta"] - ) - - R_deriv, z_deriv = flux_surfaces.convert_to_Rz(rho_arr, theta_arr, t_arr) - - R_deriv = cast(DataArray, R_deriv).transpose("rho_poloidal", "theta", "t") - z_deriv = cast(DataArray, z_deriv).transpose("rho_poloidal", "theta", "t") - - data_rho_theta = data_R_z.indica.interp2d( - {"R": R_deriv, "z": z_deriv}, method="linear", assume_sorted=True - ) - data_rho_theta = data_rho_theta.transpose("rho_poloidal", "theta", "t") - - return data_rho_theta, R_deriv, z_deriv - - def basic_extrapolation( - self, - data_rho_theta: DataArray, - electron_density: DataArray, - threshold_rho: float, - ): - """Basic extrapolation which eliminates the data_rho_theta for - rho > threshold_rho and joins on electron density from that point - outwards (in rho). Also multiplies electron density to prevent a - discontinuity at rho_poloidal=threshold_rho. - - Parameters - ---------- - data_rho_theta - xarray.DataArray to extrapolate. Dimensions (rho, theta, t) - electron_density - xarray.DataArray of Electron density (axisymmetric). Dimensions (rho, t) - threshold_rho - Threshold value (float) of rho beyond which SXR diagnostics cannot - be used to accurately infer impurity density. - - Returns - ------- - extrapolated_impurity_density - xarray.DataArray of extrapolated impurity density. - Dimensions (rho, theta, t) - """ - - theta_arr = np.array([0.0, np.pi], dtype=float) - theta_arr = DataArray( # type: ignore - data=theta_arr, coords={"theta": theta_arr}, dims=["theta"] - ) - - boundary_electron_density = electron_density.sel( - {"rho_poloidal": threshold_rho} - ).squeeze() - boundary_data = data_rho_theta.sel({"rho_poloidal": threshold_rho}).squeeze() - - discontinuity_scale = boundary_data / boundary_electron_density - - # Continue impurity_density_sxr following the shape of the electron density - # profile. - bounded_data = data_rho_theta.where( - data_rho_theta.rho_poloidal <= threshold_rho, 0.0 - ) - - bounded_electron_density = electron_density.where( - electron_density.rho_poloidal > threshold_rho, 0.0 - ) - - extrapolated_impurity_density = ( - bounded_data + bounded_electron_density * discontinuity_scale - ) - - return extrapolated_impurity_density - - def extrapolation_smoothing( - self, - extrapolated_data: DataArray, - rho_arr: DataArray, - ): - """Function to smooth extrapolatd data. Extrapolated data may not have - any 0th order discontinuity but 1st order discontinuities may exist. - Smoothing is necessary to eliminate these higher order discontinuities. - - Parameters - ---------- - extrapolated_data - xarray.DataArray extrapolated data to be smoothed. - Dimensions (rho, theta, t) - rho_arr - xarray.DataArray used to construct smoothing splines. Dimensions (rho) - (Must be higher or the same resolution as the rho dimension - of extrapolated_data) - - Returns - ------- - extrapolated_smooth_lfs_arr - Extrapolated smoothed data on low-field side (fixed theta = 0) - extrapolated_smooth_hfs_arr - Extrapolated smoothed data on high-field side (fixed theta = pi) - """ - t = extrapolated_data.coords["t"] - - extrapolated_smooth_lfs = [] - extrapolated_smooth_hfs = [] - - for ind_t, it in enumerate(extrapolated_data.coords["t"]): - variance_extrapolated_data_lfs = extrapolated_data.isel( - {"t": ind_t, "theta": 0} - ).var("rho_poloidal") - - variance_extrapolated_data_hfs = extrapolated_data.isel( - {"t": ind_t, "theta": 1} - ).var("rho_poloidal") - - extrapolated_spline_lfs = UnivariateSpline( - rho_arr, - extrapolated_data.isel(t=ind_t).sel(theta=0), - k=5, - s=0.001 * variance_extrapolated_data_lfs, - ) - - extrapolated_spline_hfs = UnivariateSpline( - rho_arr, - extrapolated_data.isel(t=ind_t).sel(theta=np.pi), - k=5, - s=0.001 * variance_extrapolated_data_hfs, - ) - - extrapolated_smooth_lfs.append(extrapolated_spline_lfs(rho_arr, 0)) - extrapolated_smooth_hfs.append(extrapolated_spline_hfs(rho_arr, 0)) - - extrapolated_smooth_lfs_arr = DataArray( - data=extrapolated_smooth_lfs, - coords={"t": t, "rho_poloidal": rho_arr}, - dims=["t", "rho_poloidal"], - ) - - extrapolated_smooth_hfs_arr = DataArray( - data=extrapolated_smooth_hfs, - coords={"t": t, "rho_poloidal": rho_arr}, - dims=["t", "rho_poloidal"], - ) - - extrapolated_smooth_lfs_arr = extrapolated_smooth_lfs_arr.transpose( - "rho_poloidal", "t" - ) - extrapolated_smooth_hfs_arr = extrapolated_smooth_hfs_arr.transpose( - "rho_poloidal", "t" - ) - - # Following section is to ensure that near the rho_poloidal=0 region, the - # extrapolated_smooth_data is constant (ie. with a first-order derivative of 0). - inv_extrapolated_smooth_hfs = DataArray( - data=np.flip(extrapolated_smooth_hfs_arr.data, axis=0), - coords={ - "rho_poloidal": -1 - * np.flip(extrapolated_smooth_hfs_arr.coords["rho_poloidal"].data), - "t": extrapolated_smooth_hfs_arr.coords["t"].data, - }, - dims=["rho_poloidal", "t"], - ) - - inv_rho_arr = inv_extrapolated_smooth_hfs.coords["rho_poloidal"].data - inv_del_val = inv_rho_arr[-1] - - inv_extrapolated_smooth_hfs = inv_extrapolated_smooth_hfs.drop_sel( - rho_poloidal=inv_del_val - ) - - extrapolated_smooth_mid_plane_arr = concat( - (inv_extrapolated_smooth_hfs, extrapolated_smooth_lfs_arr), "rho_poloidal" - ) - - rho_zero_ind = np.where( - np.isclose(extrapolated_smooth_mid_plane_arr.rho_poloidal.data, 0.0) - )[0][0] - - smooth_central_region = extrapolated_smooth_mid_plane_arr.isel( - rho_poloidal=slice(rho_zero_ind - 2, rho_zero_ind + 3) - ) - - smooth_central_region.loc[:, :] = smooth_central_region.max(dim="rho_poloidal") - - extrapolated_smooth_mid_plane_arr.loc[ - extrapolated_smooth_mid_plane_arr.rho_poloidal.data[ - rho_zero_ind - 2 - ] : extrapolated_smooth_mid_plane_arr.rho_poloidal.data[rho_zero_ind + 2], - :, - ] = smooth_central_region - - inv_extrapolated_smooth_hfs = extrapolated_smooth_mid_plane_arr.isel( - rho_poloidal=slice(0, rho_zero_ind + 1) - ) - - extrapolated_smooth_hfs_arr = DataArray( - data=np.flip(inv_extrapolated_smooth_hfs.data, axis=0), - coords=extrapolated_smooth_hfs_arr.coords, - dims=extrapolated_smooth_hfs_arr.dims, - ) - - # Ignoring mypy warning since it seems to be unaware that the xarray .loc - # method uses label-based indexing and slicing instead of integer-based. - extrapolated_smooth_lfs_arr = extrapolated_smooth_mid_plane_arr.loc[ - 0: # type: ignore - ] - - return extrapolated_smooth_lfs_arr, extrapolated_smooth_hfs_arr - - def apply_asymmetry( - self, - asymmetry_parameter: DataArray, - extrapolated_smooth_hfs: DataArray, - extrapolated_smooth_lfs: DataArray, - R_deriv: DataArray, - ): - """Applying an asymmetry parameter to low-field-side data which - will be extended over the poloidal extent to obtain an asymmetric - extrapolated smoothed data on a (rho, theta) grid. - - Parameters - ---------- - asymmetry_parameter - Asymmetry parameter to apply. - xarray.DataArray with dimensions (rho, t) - extrapolated_smooth_hfs - Extrapolated smoothed data on high-field side (fixed theta = pi). - xarray.DataArray with dimensions (rho, t) - extrapolated_smooth_lfs - Extrapolated smoothed data on low-field side (fixed theta = 0). - xarray.DataArray with dimensions (rho, t) - R_deriv - Variable describing value of R in every coordinate on a (rho, theta) grid. - xarray.DataArray with dimensions (rho, theta, t) - - Returns - ------- - extrapolated_smooth_data - Extrapolated and smoothed data on full (rho, theta) grid. - xarray.DataArray with dimensions (rho, theta, t) - asymmetry_modifier - Asymmetry modifier used to transform a low-field side only rho-profile - of a poloidally asymmetric quantity to a full poloidal cross-sectional - profile ie. (rho, t) -> (rho, theta, t). Also can be defined as: - exp(asymmetry_parameter * (R ** 2 - R_lfs ** 2)), where R is the major - radius as a function of (rho, theta, t) and R_lfs is the low-field-side - major radius as a function of (rho, t). xarray DataArray with dimensions - (rho, theta, t) - """ - rho_arr = extrapolated_smooth_hfs.coords["rho_poloidal"] - self.rho_arr = rho_arr - - theta_arr = np.linspace(-np.pi, np.pi, 21) - theta_arr = DataArray( - theta_arr, {"theta": theta_arr}, ["theta"] - ) # type: ignore - - asymmetry_modifier = asymmetry_modifier_from_parameter( - asymmetry_parameter, R_deriv - ) - - extrapolated_smooth_data = extrapolated_smooth_lfs * asymmetry_modifier - extrapolated_smooth_data = extrapolated_smooth_data.transpose( - "rho_poloidal", "theta", "t" - ) - - return extrapolated_smooth_data - - def transform_to_R_z( - self, - R_deriv: DataArray, - z_deriv: DataArray, - extrapolated_smooth_data: DataArray, - flux_surfaces: FluxSurfaceCoordinates, - ): - """Function to transform data from an (rho, theta) grid to a (R, z) grid - - Parameters - ---------- - R_deriv - Variable describing value of R in every coordinate on a (rho, theta) grid. - xarray.DataArray with dimensions (rho, theta, t) - (from derive_and_apply_asymmetry) - z_deriv - Variable describing value of z in every coordinate on a (rho, theta) grid. - xarray.DataArray with dimensions (rho, theta, t) - (from derive_and_apply_asymmetry) - extrapolated_smooth_data - Extrapolated and smoothed data to transform onto (R, z) grid. - xarray.DataArray with dimensions (rho, theta, t) - flux_surfaces - FluxSurfaceCoordinates object representing polar coordinate systems - using flux surfaces for the radial coordinate. - - Returns - ------- - extrapolated_smooth_data - Extrapolated and smoothed data on (R, z) grid. - xarray.DataArray with dimensions (R, z, t) - """ - R_arr = np.linspace(np.min(R_deriv[1:]), np.max(R_deriv[1:]), 40) - z_arr = np.linspace(np.min(z_deriv), np.max(z_deriv), 40) - - R_arr = DataArray(R_arr, {"R": R_arr}, ["R"]) # type: ignore - z_arr = DataArray(z_arr, {"z": z_arr}, ["z"]) # type: ignore - - t_arr = extrapolated_smooth_data.coords["t"] - - rho_derived, theta_derived = flux_surfaces.convert_from_Rz(R_arr, z_arr, t_arr) - rho_derived = cast(DataArray, rho_derived).transpose("R", "z", "t") - theta_derived = cast(DataArray, theta_derived).transpose("R", "z", "t") - rho_derived = abs(rho_derived) - - extrapolated_smooth_data = extrapolated_smooth_data.indica.interp2d( - {"rho_poloidal": rho_derived, "theta": theta_derived}, - method="linear", - assume_sorted=True, - ) - - extrapolated_smooth_data = extrapolated_smooth_data.fillna(0.0) - - return extrapolated_smooth_data - - def fitting_function( - self, - amplitude: float, - standard_dev: float, - position: float, - ): - """Function to construct a signal that modifies the - extrapolated smoothed impurity density. The signal is constructed - using a Gaussian profile with the three free parameters. - - Parameters - ---------- - amplitude - Amplitude of the additional signal (Gaussian amplitude) - standard_dev - Standard deviation associated with the Gaussian construction - (can be defined as FWHM/2.355 where FWHM is full-width at half maximum) - position - Position of the Gaussian. During optimization this is constrained to - the extrapolated region of rho (ie. outside the SXR validity region). - - Returns - ------- - sig - xarray.DataArray containing the Gaussian signal with dimensions (rho) - """ - rho_arr = self.rho_arr - - gaussian_signal = norm(loc=position, scale=standard_dev) - - sig = gaussian_signal.pdf(rho_arr) - - sig = DataArray( - data=sig, coords={"rho_poloidal": rho_arr}, dims=["rho_poloidal"] - ) - - sig /= sig.max() - - sig *= amplitude - - return sig - - def optimize_perturbation( - self, - extrapolated_smooth_data: DataArray, - orig_bolometry_data: DataArray, - bolometry_obj: BolometryDerivation, - impurity_element: str, - asymmetry_parameter: DataArray, - threshold_rho: DataArray, - R_deriv: DataArray, - time_correlation: bool = True, - ): - """Optimizes a Gaussian-style perturbation to recover the over-density - structure that is expected on the low-field-side of the plasma. - - Parameters - ---------- - extrapolated_smooth_data - Extrapolated and smoothed data which continues the impurity density - beyond the soft x-ray threshold limit by using electron density as a - guide. xarray DataArray with dimensions (rho, theta, t). - orig_bolometry_data - Original bolometry data that is used in the objective function to fit - the perturbation. xarray DataArray with dimensions (channels, t). - bolometry_obj - BolometryDerivation object. - impurity_element - String of impurity element symbol. - asymmetry_parameter - Asymmetry parameter used to transform a low-field side only rho-profile - of a poloidally asymmetric quantity to a full poloidal cross-sectional - profile ie. (rho, t) -> (rho, theta, t). Dimensions (rho, t) - threshold_rho - Threshold rho value beyond which asymmetry parameter is invalid and - should be fitted from bolometry. - R_deriv - Variable describing value of R in every coordinate on a (rho, theta) grid. - xarray.DataArray with dimensions (rho, theta, t) - time_correlation - Boolean to indicate whether or not to use time correlated guesses during - the optimization (ie. the result of the optimization for the previous - time-step is used as a guess for the next time-step.) - - Returns - ------- - fitted_density - New density with an optimized perturbation on the low-field-side that - matches the original bolometry data. xarray DataArray with dimensions - (rho, theta, t) - """ - - input_check( - "extrapolated_smooth_data", - extrapolated_smooth_data, - DataArray, - ndim_to_check=3, - strictly_positive=False, - ) - - input_check( - "orig_bolometry_data", - orig_bolometry_data, - DataArray, - ndim_to_check=2, - strictly_positive=False, - ) - - input_check( - "asymmetry_parameter", - asymmetry_parameter, - DataArray, - ndim_to_check=2, - positive=False, - strictly_positive=False, - ) - - rho_arr = self.rho_arr - drho = np.max(np.diff(rho_arr)) - - # Check whether bolometry_obj as been called at least once - # (ie. does it have a trimmed variant of the LoS bolometry data.) - if hasattr(bolometry_obj, "LoS_bolometry_data_trimmed"): - trim = True - orig_bolometry_trimmed = [] - for bolo_diag, bolo_los in zip( - orig_bolometry_data.coords["bolo_kb5v_coords"], - [i[6] for i in bolometry_obj.LoS_bolometry_data], - ): - LoS_bolometry_data_trimmed_labels = [ - i[6] for i in bolometry_obj.LoS_bolometry_data_trimmed - ] - if bolo_los in LoS_bolometry_data_trimmed_labels: - orig_bolometry_trimmed.append(orig_bolometry_data.loc[:, bolo_diag]) - orig_bolometry = concat( - orig_bolometry_trimmed, dim="bolo_kb5v_coords" - ).assign_attrs(**orig_bolometry_data.attrs) - else: - trim = False - orig_bolometry = orig_bolometry_data - - extrapolated_smooth_data_mean = np.mean( - extrapolated_smooth_data.loc[self.threshold_rho[0] :, :, :] - ) - - def objective_func(objective_array: Sequence, time: float): - """Objective function that is passed to scipy.optimize.least_squares - - Parameters - ---------- - objective_array - List of: - [amplitude, standard deviation, position and asymmetry_parameter] - defining the Gaussian perturbation. - time - Float specifying the time point of interest. - - Returns - ------- - abs_error - Absolute error between the derived bolometry data (with a given - Gaussian perturbation) and the original bolometry data (ground truth) - for the selected time point. - xarray.DataArray with dimensions (channels) - """ - ( - amplitude, - standard_dev, - position, - edge_asymmetry_parameter, - ) = objective_array - - perturbation_signal = self.fitting_function( - amplitude, standard_dev, position - ) - - # trim perturbation_signal to only be valid within rho = 0.0 and rho = 1.0 - perturbation_signal = perturbation_signal.interp( - rho_poloidal=rho_arr, method="linear" - ) - - # fit asymmetry parameter for region where it was invalid - asymmetry_parameter_cont = asymmetry_parameter.where( - asymmetry_parameter.rho_poloidal < threshold_rho, - other=edge_asymmetry_parameter, - ) - asymmetry_modifier = asymmetry_modifier_from_parameter( - asymmetry_parameter_cont, R_deriv - ) - - perturbation_signal = perturbation_signal * asymmetry_modifier.sel(t=time) - - modified_density = ( - extrapolated_smooth_data.sel(t=time) + perturbation_signal - ) - - bolometry_obj.impurity_densities.loc[ - impurity_element, :, :, time - ] = modified_density - - # mypy unable to determine the length of bolometry_args so is getting - # confused about whether t_val is included in bolometry_args or not, - # hence ignored - modified_bolometry_data = bolometry_obj( # type:ignore - deriv_only=True, trim=trim, t_val=time - ) - - comparison_orig_bolometry_data = orig_bolometry.sel( - t=time, method="nearest", drop=True - ) - - error = np.abs( - modified_bolometry_data.data - comparison_orig_bolometry_data.data - ) - - error = np.nan_to_num(error) - - return error - - fitted_density = zeros_like(bolometry_obj.electron_density) - - lower_amp_bound = 0.1 * extrapolated_smooth_data_mean - upper_amp_bound = 1e4 * extrapolated_smooth_data_mean - - lower_width_bound = (drho / self.halfwidthhalfmax_coeff) / 3 - upper_width_bound = ( - (1.1 - self.threshold_rho[0].data) / self.fullwidthhalfmax_coeff - ) / 3 - - lower_pos_bound = self.threshold_rho[0].data + drho - upper_pos_bound = 1.1 - - initial_guesses = np.array( - [ - [ - extrapolated_smooth_data_mean, - np.mean([lower_width_bound, upper_width_bound]), - np.mean([lower_pos_bound, upper_pos_bound]), - # asymmetry initial guess is value at threshold - asymmetry_parameter.isel(t=0).sel( - rho_poloidal=threshold_rho, method="ffill" - ), - ] - ] - ) - - fitting_bounds = [ - np.array( - [ - lower_amp_bound, - lower_width_bound, - lower_pos_bound, - -10 * np.abs(asymmetry_parameter.max()), - ] - ), - np.array( - [ - upper_amp_bound, - upper_width_bound, - upper_pos_bound, - 10 * np.abs(asymmetry_parameter.max()), - ] - ), - ] - - # set scale of optimizer steps for amp, width, position and asymmetry - x_scale = np.array( - [ - extrapolated_smooth_data_mean, - 0.3, - 1, - asymmetry_parameter.std(), - ] - ) - - result = least_squares( - fun=objective_func, - x0=initial_guesses[0], - bounds=tuple(fitting_bounds), - max_nfev=50, - args=(extrapolated_smooth_data.coords["t"].data[0],), - ftol=1e-15, - xtol=1e-60, - gtol=1e-60, - x_scale=x_scale, - ) - - gaussian_params = result.x - - if time_correlation: - initial_guesses = np.append( - initial_guesses, np.array([gaussian_params]), axis=0 - ) - - for ind_t, it in enumerate(extrapolated_smooth_data.coords["t"].data[1:]): - upper_width_bound = ( - (1.1 - self.threshold_rho[ind_t].data) / self.fullwidthhalfmax_coeff - ) / 3 - - lower_pos_bound = self.threshold_rho[ind_t].data + drho - - fitting_bounds[0][2] = lower_pos_bound - fitting_bounds[1][1] = upper_width_bound - - if time_correlation: - result = least_squares( - fun=objective_func, - x0=initial_guesses[ind_t], - bounds=tuple(fitting_bounds), - max_nfev=50, - args=(it,), - ftol=1e-60, - xtol=1e-5, - gtol=1e-60, - x_scale=x_scale, - ) - - gaussian_params = result.x - - initial_guesses = np.append( - initial_guesses, np.array([gaussian_params]), axis=0 - ) - else: - result = least_squares( - fun=objective_func, - x0=initial_guesses[0], - bounds=tuple(fitting_bounds), - max_nfev=15, - args=(it,), - ftol=1e-15, - xtol=1e-60, - gtol=1e-60, - x_scale=x_scale, - ) - - gaussian_params = result.x - - fitted_density = bolometry_obj.impurity_densities.loc[impurity_element, :, :, :] - fitted_density = fitted_density.transpose("rho_poloidal", "theta", "t") - - return fitted_density - - def __call__( # type: ignore - self, - impurity_density_sxr: DataArray, - electron_density: DataArray, - electron_temperature: DataArray, - truncation_threshold: float, - flux_surfaces: FluxSurfaceCoordinates, - asymmetry_parameter: DataArray = None, - t: DataArray = None, - ): - """Extrapolates the impurity density beyond the limits of SXR (Soft X-ray) - - Parameters - ---------- - impurity_density_sxr - xarray.DataArray of impurity density derived from soft X-ray emissivity. - Dimensions (rho_poloidal, theta, t) or (R, z, t) - electron_density - xarray.DataArray of electron density. Dimensions (rho ,t) - electron_temperature - xarray.DataArray of electron temperature. Dimensions (rho, t) - truncation_threshold - Truncation threshold (float) for the electron temperature - flux_surfaces - FluxSurfaceCoordinates object representing polar coordinate systems - using flux surfaces for the radial coordinate. - asymmetry_parameter - Optional, asymmetry parameter (either externally sourced or pre-calculated), - xarray.DataArray with dimensions (rho, t) - t - Optional float, time at which the impurity concentration is to - be calculated at. - - Returns - ------- - extrapolated_smooth_density_R_z - Extrapolated and smoothed impurity density ((R, z) grid). - xarray.DataArray with dimensions (R, z, t) - extrapolated_smooth_density_rho_theta - Extrapolated and smoothed impurity density ((rho, theta) grid). - xarray.DataArray with dimensions (rho, theta, t). - t - If ``t`` was not specified as an argument for the __call__ function, - return the time the results are given for. - Otherwise return the argument. - """ - - input_check( - "impurity_density_sxr", - impurity_density_sxr, - DataArray, - ndim_to_check=3, - strictly_positive=False, - ) - - input_check( - "electron_density", - electron_density, - DataArray, - ndim_to_check=2, - strictly_positive=False, - ) - - input_check( - "electron_temperature", - electron_temperature, - DataArray, - strictly_positive=True, - ) - - input_check( - "truncation_threshold", - truncation_threshold, - float, - strictly_positive=True, - ) - - input_check("flux_surfaces", flux_surfaces, FluxSurfaceCoordinates) - - if t is None: - t = electron_density.t - else: - input_check("t", t, DataArray, strictly_positive=False) - - self.threshold_rho = recover_threshold_rho( - truncation_threshold, electron_temperature - ) - - # Transform impurity_density_sxr to (rho, theta) coordinates - rho_arr = electron_density.coords["rho_poloidal"] - t_arr = t - if set(["R", "z"]).issubset(set(list(impurity_density_sxr.dims))): - ( - impurity_density_sxr_rho_theta, - R_deriv, - z_deriv, - ) = self.transform_to_rho_theta( - impurity_density_sxr, - flux_surfaces, - rho_arr, - t_arr=t_arr, - ) - elif set(["rho_poloidal", "theta"]).issubset( - set(list(impurity_density_sxr.dims)) - ): - impurity_density_sxr_rho_theta = impurity_density_sxr - - R_deriv, z_deriv = flux_surfaces.convert_to_Rz( - impurity_density_sxr_rho_theta.coords["rho_poloidal"], - impurity_density_sxr_rho_theta.coords["theta"], - impurity_density_sxr_rho_theta.coords["t"], - ) - else: - raise ValueError( - 'Inputted impurity_density_sxr does not have any compatible\ - coordinates: ["rho_poloidal", "theta"] or ["R", "z"]' - ) - - # Continue impurity_density_sxr following the shape of the electron density - # profile and mitigate discontinuity. - - extrapolated_impurity_density = self.basic_extrapolation( - impurity_density_sxr_rho_theta, electron_density, self.threshold_rho - ) - - assert np.all(np.logical_not(np.isnan(extrapolated_impurity_density))) - - # Smoothing extrapolated data at the discontinuity. - # (There is still a discontinuity in the radial gradient.) - - extrapolated_smooth_lfs, extrapolated_smooth_hfs = self.extrapolation_smoothing( - extrapolated_impurity_density, rho_arr - ) - - if asymmetry_parameter is None: - asymmetry_parameter = asymmetry_from_rho_theta( - impurity_density_sxr_rho_theta, flux_surfaces, self.threshold_rho, t_arr - ) - else: - input_check( - "asymmetry_parameter", - asymmetry_parameter, - DataArray, - ndim_to_check=2, - positive=False, - strictly_positive=False, - ) - - # Applying the asymmetry parameter to extrapolated density. - # Also extends the data beyond the hfs and lfs to be - # the full poloidal angle range. - - extrapolated_smooth_density_rho_theta = self.apply_asymmetry( - asymmetry_parameter, - extrapolated_smooth_hfs, - extrapolated_smooth_lfs, - R_deriv, - ) - - # Transform extrapolated density back onto a (R, z) grid - - extrapolated_smooth_density_R_z = self.transform_to_R_z( - R_deriv, z_deriv, extrapolated_smooth_density_rho_theta, flux_surfaces - ) - - asymmetry_modifier = asymmetry_modifier_from_parameter( - asymmetry_parameter, R_deriv - ) - self.asymmetry_modifier = asymmetry_modifier - - return ( - extrapolated_smooth_density_R_z, - extrapolated_smooth_density_rho_theta, - t, - ) diff --git a/indica/operators/impurity_concentration.py b/indica/operators/impurity_concentration.py deleted file mode 100644 index 6b64c207..00000000 --- a/indica/operators/impurity_concentration.py +++ /dev/null @@ -1,245 +0,0 @@ -"""Operator calculating the impurity concentration -of a given element. -""" - -from typing import List -from typing import Tuple -from typing import Union - -import numpy as np -from xarray import DataArray -from xarray.core.common import zeros_like - -from indica.converters.flux_surfaces import FluxSurfaceCoordinates -from .abstractoperator import EllipsisType -from .abstractoperator import Operator -from .. import session -from ..datatypes import DataType -from ..utilities import input_check - - -class ImpurityConcentration(Operator): - """Calculate impurity concentration of a given element. - - Attributes - ---------- - ARGUMENT_TYPES: List[DataType] - Ordered list of the types of data expected for each argument of the - operator. - RESULT_TYPES: List[DataType] - Ordered list of the types of data returned by the operator. - - Returns - ------- - concentration - xarray.DataArray containing the impurity concentration for the - given impurity element. - t - If ``t`` was not specified as an argument for the __call__ function, - return the time the results are given for. - Otherwise return the argument. - - Methods - ------- - __call__( - element, Zeff_LoS, impurity_densities, electron_density, - mean_charge, flux_surfaces, t - ) - Calculates the impurity concentration for the inputted element. - """ - - ARGUMENT_TYPES: List[Union[DataType, EllipsisType]] = [] - - RESULT_TYPES: List[Union[DataType, EllipsisType]] = [ - ("impurity_concentration", "impurity_element"), - ("time", "impurity_element"), - ] - - def __init__(self, sess: session.Session = session.global_session): - super().__init__(sess=sess) - - def return_types(self, *args: DataType) -> Tuple[DataType, ...]: - return super().return_types(*args) - - def __call__( # type: ignore - self, - element: str, - Zeff_LoS: DataArray, - impurity_densities: DataArray, - electron_density: DataArray, - mean_charge: DataArray, - flux_surfaces: FluxSurfaceCoordinates, - t: DataArray = None, - ): - """Calculates the impurity concentration for the inputted element. - - Parameters - ---------- - element - String specifying the symbol of the element for which the impurity - concentration is desired. - Zeff_LoS - xarray.DataArray containing the Zeff value/s from Bremsstrahlung (ZEFH/KS3) - impurity_densities - xarray.DataArray of impurity densities for all impurity elements - of interest. - electron_density - xarray.DataArray of electron density - mean_charge - xarray.DataArray of mean charge of all impurity elements of interest. - This can be provided manually (with dimensions of - ["element", "rho_poloidal", "t]), - or can be passed as the results of MeanCharge.__call__ - flux_surfaces - FluxSurfaceCoordinates object that defines the flux surface geometry - of the equilibrium of interest. - t - Optional, time at which the impurity concentration is to be calculated at. - - Returns - ------- - concentration - xarray.DataArray containing the impurity concentration for the - given impurity element. - t - If ``t`` was not specified as an argument for the __call__ function, - return the time the results are given for. - Otherwise return the argument. - """ - input_check( - "impurity_densities", - impurity_densities, - DataArray, - strictly_positive=False, - ) - - input_check("element", element, str) - - elements_list = impurity_densities.coords["element"] - - try: - assert element in elements_list - except AssertionError: - raise ValueError( - f"Please input a single valid element from list:\ - {elements_list}" - ) - - if t is None: - t = Zeff_LoS.t - else: - input_check("t", t, DataArray, ndim_to_check=1, strictly_positive=False) - - input_check( - "Zeff_LoS", - Zeff_LoS, - DataArray, - ndim_to_check=1, - strictly_positive=False, - ) - input_check( - "electron_density", - electron_density, - DataArray, - ndim_to_check=2, - strictly_positive=True, - ) - input_check( - "mean_charge", - mean_charge, - DataArray, - ndim_to_check=3, - strictly_positive=False, - ) - input_check("flux_surfaces", flux_surfaces, FluxSurfaceCoordinates) - - Zeff_LoS = Zeff_LoS.interp(t=t, method="nearest") - - transform = Zeff_LoS.attrs["transform"] - x1_name = transform.x1_name - x2_name = transform.x2_name - - x1 = Zeff_LoS.coords[x1_name] - x2_arr = np.linspace(0, 1, 300) - x2 = DataArray(data=x2_arr, dims=[x2_name]) - - R_arr, z_arr = transform.convert_to_Rz(x1, x2, t) - - rho, theta = flux_surfaces.convert_from_Rz(R_arr, z_arr, t) - - if isinstance(R_arr, (DataArray, np.ndarray)): - R_arr = R_arr.squeeze() - - if isinstance(rho, (DataArray, np.ndarray)): - rho = rho.squeeze() - if isinstance(rho, DataArray): - rho = rho.drop_vars("t") - rho = rho.drop_vars("R") - rho = rho.drop_vars("z") - rho = rho.fillna(2.0) - - if set(["R", "z"]).issubset(set(list(impurity_densities.dims))): - impurity_densities = impurity_densities.indica.interp2d( - z=z_arr, - R=R_arr, - method="cubic", - assume_sorted=True, - ) - elif set(["rho_poloidal", "theta"]).issubset( - set(list(impurity_densities.dims)) - ): - impurity_densities = impurity_densities.interp( - rho_poloidal=rho, theta=theta, method="linear", assume_sorted=True - ) - elif set(["rho_poloidal"]).issubset(set(list(impurity_densities.dims))): - impurity_densities = impurity_densities.interp( - rho_poloidal=rho, method="linear", assume_sorted=True - ) - else: - raise ValueError( - 'Inputted impurity densities does not have any compatible\ - coordinates: ["rho_poloidal", "theta"], ["rho_poloidal"]\ - or ["R", "z"]' - ) - - impurity_densities = impurity_densities.interp( - t=t, method="linear", assume_sorted=True - ) - - electron_density = electron_density.interp( - rho_poloidal=rho, method="linear", assume_sorted=True - ) - - electron_density = electron_density.interp( - t=t, method="linear", assume_sorted=True - ) - - mean_charge = mean_charge.interp( - rho_poloidal=rho, method="linear", assume_sorted=True - ) - - mean_charge = mean_charge.interp(t=t, method="linear", assume_sorted=True) - - dl = transform.distance(x2_name, DataArray(0), x2[0:2], 0) - dl = dl[1] - LoS_length = dl * 300 - - concentration = zeros_like(Zeff_LoS) - - term_1 = LoS_length * (Zeff_LoS - 1) - - term_2 = zeros_like(term_1) - for k, kdens in enumerate(impurity_densities.coords["element"]): - if element == kdens: - term_3 = (mean_charge[k] ** 2 - mean_charge[k]).sum(x2_name) * dl - continue - - term2_integrand = (impurity_densities[k] / electron_density) * ( - mean_charge[k] ** 2 - mean_charge[k] - ) - - term_2 += term2_integrand.sum(x2_name) * dl - - concentration = (term_1 - term_2) / term_3 - - return concentration, t diff --git a/indica/operators/spline_fit.py b/indica/operators/spline_fit.py deleted file mode 100644 index 4f42287d..00000000 --- a/indica/operators/spline_fit.py +++ /dev/null @@ -1,322 +0,0 @@ -"""Fit 1-D splines to data.""" - -from itertools import accumulate -from itertools import chain -from itertools import repeat -from itertools import tee -from typing import cast -from typing import Dict -from typing import Hashable -from typing import List -from typing import Tuple -from typing import Union -import warnings - -import numpy as np -from scipy.interpolate import CubicSpline -from scipy.optimize import least_squares -from scipy.sparse import lil_matrix -from xarray import DataArray - -from .abstractoperator import EllipsisType -from .abstractoperator import Operator -from .. import session -from ..converters import bin_to_time_labels -from ..converters import CoordinateTransform -from ..converters import FluxSurfaceCoordinates -from ..datatypes import DataType -from ..numpy_typing import ArrayLike -from ..utilities import broadcast_spline -from ..utilities import coord_array - - -SingleBoundaryType = Union[str, Tuple[int, ArrayLike]] -BoundaryType = Union[str, Tuple[SingleBoundaryType, SingleBoundaryType]] - - -class Spline: - """Callable class wrapping a `:class:scipy.interpolate.CubicSpline` - object so it will work with DataArrays. It performs interpolation - over one dimension, but can do this onto a multidimensional grid. - - Parameters - ---------- - values : DataArray - The values to interpolate. - dim : Hashable - The axis along which to interpolate. - coord_transform : CoordinateTransform - The transform describing the coordinate system used by `values`. - bounds : BoundaryType - The boundary condition to pass to `:class:scipy.interpolate.CubicSpline`. - - """ - - def __init__( - self, - values: DataArray, - dim: Hashable, - coord_transform: CoordinateTransform, - bounds: BoundaryType = "clamped", - ): - self.dim = dim - self.spline_dims = tuple(d for d in values.dims if d != dim) - self.spline_coords: Dict[Hashable, np.ndarray] = { - k: np.asarray(v) for k, v in values.coords.items() if k != self.dim - } - transpose_order = (self.dim,) + self.spline_dims - self.spline = CubicSpline( - values.coords[dim], values.transpose(*transpose_order), 0, bounds, False - ) - self.transform = coord_transform - - def __call__( - self, - coord_system: CoordinateTransform, - x1: DataArray, - x2: DataArray, - t: DataArray, - ) -> DataArray: - """Get the spline values at the locations given by the - coordinates. Although it takes multiple coordinates as - arguments, the actual interpolation will only be done along - the `dim` specified at instantiation. - - Parameters - ---------- - coord_system - The transform describing the system used by the provided coordinates - x1 - The first spatial coordinate - x2 - The second spatial coordinate - t - The time coordinate - - """ - self_x1, self_x2 = cast( - Tuple[DataArray, DataArray], - coord_system.convert_to(self.transform, x1, x2, t), - ) - coord = self_x1 if self.dim == self.transform.x1_name else self_x2 - result = broadcast_spline( - self.spline, self.spline_dims, self.spline_coords, coord - ) - result.attrs["transform"] = coord_system - return result - - -class SplineFit(Operator): - """Fit a 1-D spline to data. The spline will be given on poloidal flux - surface coordinates, as specified by the user. It can derive a - single spline fit for multiple DataArray arguments simultaneously. - - Parameters - ---------- - knots : ArrayLike - A 1-D array containing the location of spline knots to use when - fitting the data. - lower_bound : ArrayLike - The lower bounds to use for values at each not. May be either a - scalar or an array of the same shape as ``knots``. - upper_bound : ArrayLike - The upper bounds to use for values at each not. May be either a - scalar or an array of the same shape as ``knots``. - sess : session.Session - An object representing the session being run. Contains information - such as provenance data. - - """ - - ARGUMENT_TYPES: List[Union[DataType, EllipsisType]] = [ - ("norm_flux_pol", "plasma"), - ("time", "plasma"), - ("temperature", "electrons"), - ..., - ] - - def __init__( - self, - knots: ArrayLike = [0.0, 0.3, 0.6, 0.85, 0.95, 1.05], - lower_bound: ArrayLike = -np.inf, - upper_bound: ArrayLike = np.inf, - sess: session.Session = session.global_session, - ): - self.knots = coord_array(knots, "rho_poloidal") - self.lower_bound = lower_bound - if isinstance(lower_bound, np.ndarray) and lower_bound.size != self.knots.size: - raise ValueError( - "lower_bound must be either a scalar or array of same size as knots" - ) - self.upper_bound = upper_bound - if isinstance(upper_bound, np.ndarray) and upper_bound.size != self.knots.size: - raise ValueError( - "lower_bound must be either a scalar or array of same size as knots" - ) - self.spline: Spline - self.spline_vals: DataArray - super().__init__( - sess, - knots=str(knots), - lower_bound=str(lower_bound), - upper_bound=str(upper_bound), - ) - - def return_types(self, *args: DataType) -> Tuple[DataType, ...]: - """Indicates the datatypes of the results when calling the operator - with arguments of the given types. It is assumed that the - argument types are valid. - - Parameters - ---------- - args - The datatypes of the parameters which the operator is to be called with. - - Returns - ------- - : - The datatype of each result that will be returned if the operator is - called with these arguments. - - """ - input_type = args[-1] - return (input_type,) * len(args) - - def __call__( # type: ignore[override] - self, - rho: DataArray, - times: DataArray, - *data: DataArray, - ) -> Tuple[DataArray, ...]: - """Fit a spline to the provided data. - - Parameters - ---------- - rho - The poloidal flux values on which to return the result. - times - The times at which to bin the data and return the result. - data - The data to fit the spline to. - - Returns - ------- - : - The results of the fit on the give \\rho and time values. - It contains the attribute `splines` which can be used to - interpolate results onto arbitrary coordinates. - - """ - self.validate_arguments(rho, times, *data) - n_knots = len(self.knots) - flux_surfaces = FluxSurfaceCoordinates("poloidal") - flux_surfaces.set_equilibrium(data[0].indica.equilibrium) - binned_data = [bin_to_time_labels(times.data, d) for d in data] - droppable_dims = [ - [dim for dim in d.dims if dim != d.attrs["transform"].x1_name] for d in data - ] - good_channels: List[np.ndarray] = [ - np.ravel( - cast( - DataArray, - np.logical_not(np.isnan(d.isel({dim: 0 for dim in droppable}))), - ).drop_vars(droppable) - ) - for d, droppable in zip(data, droppable_dims) - ] - for d, g in zip(binned_data, good_channels): - d.attrs["nchannels"] = ( - d.size - * int(np.sum(g)) - // (d.coords[d.attrs["transform"].x1_name].size * times.size) - ) - nt = len(times) - rows = sum(d.attrs["nchannels"] for d in binned_data) * nt - cols = (n_knots - 1) * nt - sparsity = lil_matrix((rows, cols), dtype=int) - nc1, nc2 = tee(d.attrs["nchannels"] for d in binned_data) - - for nc, data_row_start in zip( - nc1, accumulate(map(lambda x: x * nt, chain(repeat(0, 1), nc2))) - ): - for i in range(nt): - rstart = data_row_start + i * nc - rend = rstart + nc - cstart = i * (n_knots - 1) - cend = cstart + (n_knots - 1) - sparsity[rstart:rend, cstart:cend] = 1 - - def knotvals_to_xarray(knotvals): - all_knots = np.empty((nt, n_knots)) - all_knots[:, :-1] = knotvals.reshape(nt, n_knots - 1) - all_knots[:, -1] = 0.0 - return DataArray( - all_knots, coords=[("t", times), ("rho_poloidal", self.knots)] - ) - - # TODO: Consider how to handle locations outside of interpolation range. - # For now just setting the interpolated values to 0.0 - def residuals(knotvals): - self.spline_vals = knotvals_to_xarray(knotvals) - self.spline = Spline(self.spline_vals, "rho_poloidal", flux_surfaces) - start = 0 - resid = np.empty(rows) - for d, g in zip(binned_data, good_channels): - end = start + d.attrs["nchannels"] * nt - rho, theta = d.indica.convert_coords(flux_surfaces) - temp_resid = ( - self.spline(flux_surfaces, rho, theta, times).fillna(0.0) - d - ).isel({d.attrs["transform"].x1_name: g}) - if d.ndim == 2: - resid[start:end] = np.ravel( - temp_resid.transpose("t", d.attrs["transform"].x1_name) - ) - elif d.ndim == 3: - resid[start:end] = np.ravel( - temp_resid.transpose( - "t", - d.attrs["transform"].x1_name, - d.attrs["transform"].x2_name, - ) - ) - start = end - # assert np.all(np.isfinite(resid)) - return resid - - guess = np.concatenate( - tuple( - np.mean([d.sel(t=t).mean() for d in binned_data]) * np.ones(n_knots - 1) - for t in times - ) - ) - - fit = least_squares( - residuals, - guess, - bounds=(self.lower_bound, self.upper_bound), - jac_sparsity=sparsity, - verbose=2, - ) - - if fit.status == -1: - raise RuntimeError( - "Improper input to `least_squares` function when trying to " - "fit emissivity to radiation data." - ) - elif fit.status == 0: - warnings.warn( - "Attempt to fit splines reached maximum number of function " - "evaluations.", - RuntimeWarning, - ) - result = self.spline(flux_surfaces, rho, DataArray(0.0), times) - result.attrs["splines"] = self.spline - self.spline_vals.attrs["datatype"] = result.attrs["datatype"] = data[0].attrs[ - "datatype" - ] - self.spline_vals.name = "knot_values" - self.assign_provenance(result) - self.assign_provenance(self.spline_vals) - for d in binned_data: - self.assign_provenance(d) - return result, self.spline_vals, *binned_data diff --git a/indica/operators/spline_fit_easy.py b/indica/operators/spline_fit_easy.py index d6549067..d514abd3 100644 --- a/indica/operators/spline_fit_easy.py +++ b/indica/operators/spline_fit_easy.py @@ -98,7 +98,7 @@ def residuals(yknots): return yspl -def spline_fit_ts( +def example_run( pulse: int = 11314, tstart: float = 0.03, tend: float = 0.1, @@ -180,5 +180,5 @@ def spline_fit_ts( if __name__ == "__main__": plt.ioff() - spline_fit_ts(11089, quantity="ne") + example_run(11089, quantity="ne") plt.show() diff --git a/indica/readers/manage_data.py b/indica/readers/manage_data.py deleted file mode 100644 index a35c029e..00000000 --- a/indica/readers/manage_data.py +++ /dev/null @@ -1,204 +0,0 @@ -from copy import deepcopy - -import numpy as np -import xarray as xr -from xarray import DataArray - -from indica.converters import FluxSurfaceCoordinates -from indica.converters.time import convert_in_time_dt - - -def initialize_bckc(data): - """ - Initialise back-calculated data dictionary of dictionaries containing all info from - the original data, apart from provenance and revision attributes - - Parameters - ---------- - data - Dictionary of diagnostics and quantities as returned by manage_data - - Returns - ------- - bckc - New dictionary initialized to nans and with limited attributes - - """ - bckc = {} - for diagnostic in data.keys(): - bckc[diagnostic] = {} - for quantity in data[diagnostic].keys(): - bckc[diagnostic][quantity] = initialize_bckc_dataarray( - data[diagnostic][quantity] - ) - - return bckc - - -def initialize_bckc_dataarray(dataarray: DataArray): - bckc_data = xr.full_like(dataarray, np.nan) - attrs = bckc_data.attrs - if type(bckc_data) == DataArray: - if "error" in attrs.keys(): - attrs["error"] = xr.full_like(attrs["error"], np.nan) - if "partial_provenance" in attrs.keys(): - attrs.pop("partial_provenance") - attrs.pop("provenance") - bckc_data.attrs = attrs - - return bckc_data - - -def bin_data_in_time( - exp_data: dict, - tstart: float, - tend: float, - dt: float, -): - """ - Bin raw experimental data on the desired time axis, assign equilibrium to - transform objects - - Parameters - ---------- - raw_data - Experimental data as returned by Indica's abstractreader.py - tstart - Start of time range for which to get data. - tend - End of time range for which to get data. - dt - Time binning/interpolation - - Returns - ------- - Data dictionary identical to input value but binned on new time axis - """ - binned_data = {} - for quant in exp_data.keys(): - data = deepcopy(exp_data[quant]) - attrs = data.attrs - if "t" in data.coords: - data = convert_in_time_dt(tstart, tend, dt, data) - if "provenance" in attrs: - data.attrs["provenance"] = attrs["provenance"] - - binned_data[quant] = data - - return binned_data - - -def map_on_equilibrium( - diagnostic_data: dict, - flux_transform: FluxSurfaceCoordinates, -): - """ - Assign equilibrium and transform, map viewing LOS - - Parameters - ---------- - diagnostic_data - Experimental data of a specific instrument (see abstractreader.py) - flux_transform - Indica's FluxSurfaceTransform object - - Returns - ------- - Data dictionary identical to input with transforms set for remapping - and remapping - """ - - data = diagnostic_data[list(diagnostic_data)[0]] - if "transform" not in data.attrs: - return data - - transform = data.attrs["transform"] - if hasattr(flux_transform, "equilibrium"): - transform.set_equilibrium(flux_transform.equilibrium, force=True) - if "LineOfSightTransform" in str(data.attrs["transform"]): - transform.set_flux_transform(flux_transform) - transform.convert_to_rho_theta(t=data.t) - - for quantity in diagnostic_data.keys(): - diagnostic_data[quantity].attrs["transform"] = transform - - return diagnostic_data - - -def assign_flux_transform( - diagnostic_data: dict, - flux_transform: FluxSurfaceCoordinates, -): - """ - Assign transform necessary to map viewing LOS from Cartesian to Flux space - - Parameters - ---------- - diagnostic_data - Experimental data of a specific instrument (see abstractreader.py) - flux_transform - Indica's FluxSurfaceTransform object - - Returns - ------- - Data dictionary identical to input with transforms set for remapping - and remapping - """ - - data = diagnostic_data[list(diagnostic_data)[0]] - if "transform" not in data.attrs: - return data - - transform = data.attrs["transform"] - if hasattr(flux_transform, "equilibrium"): - transform.set_equilibrium(flux_transform.equilibrium, force=True) - if "LineOfSightTransform" in str(data.attrs["transform"]): - transform.set_flux_transform(flux_transform) - transform.convert_to_rho_theta(t=data.t) - - for quantity in diagnostic_data.keys(): - diagnostic_data[quantity].attrs["transform"] = transform - - return diagnostic_data - - -def apply_limits( - data, - diagnostic: str, - quantity=None, - val_lim=(np.nan, np.nan), - err_lim=(np.nan, np.nan), -): - """ - Set to Nan all data whose value or relative error aren't within specified limits - """ - - if quantity is None: - quantity = list(data[diagnostic]) - else: - quantity = list(quantity) - - for q in quantity: - error = None - value = data[diagnostic][q] - if "error" in value.attrs.keys(): - error = data[diagnostic][q].attrs["error"] - - if np.isfinite(val_lim[0]): - print(val_lim[0]) - value = xr.where(value >= val_lim[0], value, np.nan) - if np.isfinite(val_lim[1]): - print(val_lim[1]) - value = xr.where(value <= val_lim[1], value, np.nan) - - if error is not None: - if np.isfinite(err_lim[0]): - print(err_lim[0]) - value = xr.where((error / value) >= err_lim[0], value, np.nan) - if np.isfinite(err_lim[1]): - print(err_lim[1]) - value = xr.where((error / value) <= err_lim[1], value, np.nan) - - data[diagnostic][q].values = value.values - - return data diff --git a/snippets/test_spline_fit.py b/snippets/test_spline_fit.py deleted file mode 100644 index 64afe455..00000000 --- a/snippets/test_spline_fit.py +++ /dev/null @@ -1,46 +0,0 @@ -import getpass -import itertools - -import matplotlib.pyplot as plt -import numpy as np - -from indica.equilibrium import Equilibrium -from indica.operators import SplineFit -from indica.readers import PPFReader -from indica.utilities import coord_array - - -rho = coord_array(np.linspace(0, 1.04, 100), "rho_poloidal") -t = coord_array(np.linspace(45, 50, 11), "t") - -reader = PPFReader(90279, 45.0, 50.0) -if reader.requires_authentication: - user = input("JET username: ") - password = getpass.getpass("JET password: ") - assert reader.authenticate(user, password) - -equilib_dat = reader.get("jetppf", "efit", 0) -hrts = reader.get("jetppf", "hrts", 0) -lidr = reader.get("jetppf", "lidr", 0) - -equilibrium = Equilibrium(equilib_dat) - -# ********************************************************************* - -for data in itertools.chain(hrts.values(), lidr.values()): - data.indica.equilibrium = equilibrium - -fitter = SplineFit(lower_bound=0.0) - -Te = [lidr["te"], hrts["te"]] -results = fitter(rho, t, *Te) -Te_smoothed = results[0] -Te_spline = Te_smoothed.attrs["splines"] - -for time in t: - print(f"Time = {float(time)}") - Te_smoothed.sel(t=time).plot(label="Spline fit") - for d in results[1:]: - d.sel(t=time).plot.line("o", x="rho_poloidal", label=str(d.name) + " data") - plt.legend() - plt.show() diff --git a/tests/unit/converters/test_flux_surfaces.py b/tests/unit/converters/test_flux_surfaces.py deleted file mode 100644 index 97f6623d..00000000 --- a/tests/unit/converters/test_flux_surfaces.py +++ /dev/null @@ -1,97 +0,0 @@ -"""Tests transforms to/from coordinate systems based on flux surfaces.""" - -from unittest.mock import Mock - -from hypothesis import given -from hypothesis.strategies import composite -from hypothesis.strategies import floats -import numpy as np - -from indica.converters import FluxSurfaceCoordinates -from indica.equilibrium import Equilibrium -from indica.utilities import coord_array -from ..fake_equilibrium import fake_equilibria -from ..fake_equilibrium import flux_types -from ..strategies import arbitrary_coordinates -from ..strategies import basis_coordinates - - -@composite -def flux_coordinates(draw, domain=((0.0, 1.0), (0.0, 1.0), (0.0, 1.0))): - """Generates :py:class:`indica.converters.FluxSurfaceCoordinates` objects.""" - result = FluxSurfaceCoordinates(draw(flux_types())) - Rmag = draw(floats(*domain[0])) - if abs(Rmag) < 1e-10: - sign = 1 if Rmag == 0.0 else np.sign(Rmag) - Rmag += sign * 0.1 * (domain[0][1] - domain[0][0]) - zmag = draw(floats(*domain[1])) - result.set_equilibrium(draw(fake_equilibria(Rmag, zmag))) - return result - - -@composite -def flux_coordinates_and_axes( - draw, domain=((0.0, 1.0), (0.0, 1.0), (0.0, 1.0)), min_side=2, max_side=12 -): - """Generates :py:class:`indica.converters.FluxSurfaceCoordinates` objects, - along with axes. - - Parameters - ---------- - min_side : integer - The minimum number of elements in an unaligned dimension for the - default coordinate arrays. (Not available for all coordinate systems.) - max_side : integer - The maximum number of elements in an unaligned dimension for the - default coordinate arrays. (Not available for all coordinate systems.) - - """ - transform = draw(flux_coordinates(domain)) - x1_vals, x2_vals, t_vals = map( - np.ravel, - draw( - basis_coordinates( - (0.0, 0.0, 0.0), (1.0, 2 * np.pi, 120.0), min_side, max_side - ) - ), - ) - return ( - transform, - coord_array(x1_vals, transform.x1_name), - coord_array(x2_vals, transform.x2_name), - coord_array(t_vals, "t"), - ) - - -@given( - flux_types(), - arbitrary_coordinates(xarray=True), - arbitrary_coordinates((0.0, 0.0, None), (1.0, 2 * np.pi, None), xarray=True), -) -def test_flux_from_Rz_mock(kind, coords, expected_result): - """Test transform of data to flux coordinates.""" - equilib = Mock(spec=Equilibrium) - equilib.flux_coords.return_value = expected_result - transform = FluxSurfaceCoordinates(kind) - transform.set_equilibrium(equilib) - result = transform.convert_from_Rz(*coords) - equilib.flux_coords.assert_called_with(*coords, kind) - assert result[0] is expected_result[0] - assert result[1] is expected_result[1] - - -@given( - flux_types(), - arbitrary_coordinates((0.0, 0.0, None), (1.0, 2 * np.pi, None), xarray=True), - arbitrary_coordinates(xarray=True), -) -def test_flux_to_Rz_mock(kind, coords, expected_result): - """Test transform of data from flux coordinates.""" - equilib = Mock(spec=Equilibrium) - equilib.spatial_coords.return_value = expected_result - transform = FluxSurfaceCoordinates(kind) - transform.set_equilibrium(equilib) - result = transform.convert_to_Rz(*coords) - equilib.spatial_coords.assert_called_with(*coords, kind) - assert result[0] is expected_result[0] - assert result[1] is expected_result[1] diff --git a/tests/unit/data_strategies.py b/tests/unit/data_strategies.py index 1bd6f09b..5a07d362 100644 --- a/tests/unit/data_strategies.py +++ b/tests/unit/data_strategies.py @@ -21,10 +21,7 @@ import numpy as np from xarray import DataArray -from indica.converters import FluxSurfaceCoordinates -from indica.converters import TrivialTransform import indica.datatypes as dt -from .strategies import monotonic_series from .strategies import separable_functions from .strategies import smooth_functions @@ -168,203 +165,6 @@ def incompatible_dataset_types(draw, datatype): return specific_type, general_types -@composite -def equilibrium_data( - draw, - machine_dims=((1.83, 3.9), (-1.75, 2.0)), - min_spatial_points=3, - max_spatial_points=12, - min_time_points=2, - max_time_points=15, - start_time=75.0, - end_time=80.0, - Btot_factor=None, -): - """Returns a dictionary containing the data necessary to construct an - :py:class:`indica.equilibrium.Equilibrium` object. - - Parameters - ---------- - machine_dims - The size of the reactor, ((Rmin, Rmax), (zmin, zmax)) - min_spatial_points - The minimum number of points to use for spatial coordinate axes - max_spatial_points - The maximum number of points to use for spatial coordinate axes - min_time_points - The minimum number of points to use for the time axis - max_time_points - The maximum number of points to use for the time axis - Btot_factor - If present, the equilibrium will have total magnetic field strength - Btot_factor/R. - """ - result = {} - nspace = draw(integers(min_spatial_points, max_spatial_points)) - ntime = draw(integers(min_time_points, max_time_points)) - times = np.linspace(start_time - 0.5, end_time + 0.5, ntime) - tfuncs = smooth_functions((start_time, end_time), 0.01) - r_centre = (machine_dims[0][0] + machine_dims[0][1]) / 2 - z_centre = (machine_dims[1][0] + machine_dims[1][1]) / 2 - raw_result = {} - attrs = { - "transform": TrivialTransform(), - "provenance": MagicMock(), - "partial_provenance": MagicMock(), - } - result["rmag"] = DataArray( - r_centre + draw(tfuncs)(times), coords=[("t", times)], name="rmag", attrs=attrs - ) - result["rmag"].attrs["datatype"] = ("major_rad", "mag_axis") - result["zmag"] = DataArray( - z_centre + draw(tfuncs)(times), coords=[("t", times)], name="zmag", attrs=attrs - ) - result["zmag"].attrs["datatype"] = ("z", "mag_axis") - fmin = draw(floats(0.0, 1.0)) - result["faxs"] = DataArray( - fmin + np.abs(draw(tfuncs)(times)), - {"t": times, "R": result["rmag"], "z": result["zmag"]}, - ["t"], - name="faxs", - attrs=attrs, - ) - result["faxs"].attrs["datatype"] = ("magnetic_flux", "mag_axis") - a_coeff = DataArray( - np.vectorize(lambda x: draw(floats(0.75 * x, x)))( - np.minimum( - np.abs(machine_dims[0][0] - result["rmag"]), - np.abs(machine_dims[0][1] - result["rmag"]), - ) - ), - coords=[("t", times)], - ) - if Btot_factor is None: - b_coeff = DataArray( - np.vectorize(lambda x: draw(floats(0.75 * x, x)))( - np.minimum( - np.abs(machine_dims[1][0] - result["zmag"].data), - np.abs(machine_dims[1][1] - result["zmag"].data), - ), - ), - coords=[("t", times)], - ) - n_exp = 0.5 # 1 # draw(floats(-0.5, 1.0)) - fmax = draw(floats(max(1.0, 2 * fmin), 10.0)) - result["fbnd"] = DataArray( - fmax - np.abs(draw(tfuncs)(times)), - coords=[("t", times)], - name="fbnd", - attrs=attrs, - ) - else: - b_coeff = a_coeff - n_exp = 1 - fdiff_max = Btot_factor * a_coeff - result["fbnd"] = DataArray( - np.vectorize(lambda axs, diff: axs + draw(floats(0.001 * diff, diff)))( - result["faxs"], fdiff_max.values - ), - coords=[("t", times)], - name="fbnd", - attrs=attrs, - ) - result["fbnd"].attrs["datatype"] = ("magnetic_flux", "separatrix") - thetas = DataArray( - np.linspace(0.0, 2 * np.pi, nspace, endpoint=False), dims=["arbitrary_index"] - ) - result["rbnd"] = ( - result["rmag"] - + a_coeff * b_coeff / np.sqrt(a_coeff**2 * np.tan(thetas) ** 2 + b_coeff**2) - ).assign_attrs(**attrs) - result["rbnd"].name = "rbnd" - result["rbnd"].attrs["datatype"] = ("major_rad", "separatrix") - result["zbnd"] = ( - result["zmag"] - + a_coeff - * b_coeff - / np.sqrt(a_coeff**2 + b_coeff**2 * np.tan(thetas) ** -2) - ).assign_attrs(**attrs) - result["zbnd"].name = "zbnd" - result["zbnd"].attrs["datatype"] = ("z", "separatrix") - - r = np.linspace(machine_dims[0][0], machine_dims[0][1], nspace) - z = np.linspace(machine_dims[1][0], machine_dims[1][1], nspace) - rgrid = DataArray(r, coords=[("R", r)]) - zgrid = DataArray(z, coords=[("z", z)]) - psin = ( - (-result["zmag"] + zgrid) ** 2 / b_coeff**2 - + (-result["rmag"] + rgrid) ** 2 / a_coeff**2 - ) ** (0.5 / n_exp) - psi = psin * (result["fbnd"] - result["faxs"]) + result["faxs"] - psi.name = "psi" - psi.attrs["transform"] = attrs["transform"] - psi.attrs["provenance"] = MagicMock() - psi.attrs["partial_provenance"] = MagicMock() - psi.attrs["datatype"] = ("magnetic_flux", "plasma") - result["psi"] = psi - - psin_coords = np.linspace(0.0, 1.0, nspace) - rho = np.sqrt(psin_coords) - psin_data = DataArray(psin_coords, coords=[("rho_poloidal", rho)]) - attrs["transform"] = FluxSurfaceCoordinates( - "poloidal", - ) - ftor_min = draw(floats(0.0, 1.0)) - ftor_max = draw(floats(max(1.0, 2 * fmin), 10.0)) - result["ftor"] = DataArray( - np.outer( - 1 + draw(tfuncs)(times), draw(monotonic_series(ftor_min, ftor_max, nspace)) - ), - coords=[("t", times), ("rho_poloidal", rho)], - name="ftor", - attrs=attrs, - ) - result["ftor"].attrs["datatype"] = ("toroidal_flux", "plasma") - if Btot_factor is None: - f_min = draw(floats(0.0, 1.0)) - f_max = draw(floats(max(1.0, 2 * fmin), 10.0)) - time_vals = draw(tfuncs)(times) - space_vals = draw(monotonic_series(f_min, f_max, nspace)) - f_raw = np.outer(abs(1 + time_vals), space_vals) - else: - f_raw = np.outer( - np.sqrt( - Btot_factor**2 - - (raw_result["fbnd"] - raw_result["faxs"]) ** 2 / a_coeff**2 - ), - np.ones_like(rho), - ) - f_raw[:, 0] = Btot_factor - result["f"] = DataArray( - f_raw, coords=[("t", times), ("rho_poloidal", rho)], name="f", attrs=attrs - ) - result["f"].attrs["datatype"] = ("f_value", "plasma") - result["rmjo"] = (result["rmag"] + a_coeff * psin_data**n_exp).assign_attrs( - **attrs - ) - result["rmjo"].name = "rmjo" - result["rmjo"].attrs["datatype"] = ("major_rad", "lfs") - result["rmjo"].coords["z"] = result["zmag"] - result["rmji"] = (result["rmag"] - a_coeff * psin_data**n_exp).assign_attrs( - **attrs - ) - result["rmji"].name = "rmji" - result["rmji"].attrs["datatype"] = ("major_rad", "hfs") - result["rmji"].coords["z"] = result["zmag"] - result["vjac"] = ( - 4 - * n_exp - * np.pi**2 - * result["rmag"] - * a_coeff - * b_coeff - * psin_data ** (2 * n_exp - 1) - ).assign_attrs(**attrs) - result["vjac"].name = "vjac" - result["vjac"].attrs["datatype"] = ("volume_jacobian", "plasma") - return result - - @composite def adf11_data( draw, From c5b7922841afc36a0e1e713a49debe4f46bc0695 Mon Sep 17 00:00:00 2001 From: Marco Sertoli Date: Fri, 27 Oct 2023 15:29:39 +0100 Subject: [PATCH 05/17] Renamed test for abstract reader --- tests/unit/readers/test_abstract_reader.py | 507 +++++++++++++++++++++ 1 file changed, 507 insertions(+) create mode 100644 tests/unit/readers/test_abstract_reader.py diff --git a/tests/unit/readers/test_abstract_reader.py b/tests/unit/readers/test_abstract_reader.py new file mode 100644 index 00000000..cccd6aba --- /dev/null +++ b/tests/unit/readers/test_abstract_reader.py @@ -0,0 +1,507 @@ +"""Test methods present on the base class DataReader.""" + +from copy import deepcopy +from numbers import Number +from typing import Any +from typing import Dict +from typing import List +from typing import Set +from typing import Tuple + +import numpy as np + +from indica import session +from indica.numpy_typing import RevisionLike +from indica.readers import DataReader +from indica.readers.available_quantities import AVAILABLE_QUANTITIES + + +# TODO these values should come from the machine dimensions variable of the reader + +TSTART = 0 +TEND = 10 + + +class Reader(DataReader): + """Class to read fake data""" + + MACHINE_DIMS = ((1.83, 3.9), (-1.75, 2.0)) + INSTRUMENT_METHODS = { + "thomson_scattering": "get_thomson_scattering", + "equilibrium": "get_equilibrium", + "cyclotron_emissions": "get_cyclotron_emissions", + "charge_exchange": "get_charge_exchange", + "spectrometer": "get_spectrometer", + "bremsstrahlung_spectroscopy": "get_bremsstrahlung_spectroscopy", + "radiation": "get_radiation", + "helike_spectroscopy": "get_helike_spectroscopy", + "interferometry": "get_interferometry", + "diode_filters": "get_diode_filters", + } + + def __init__( + self, + pulse: int, + tstart: float, + tend: float, + server: str = "", + default_error: float = 0.05, + max_freq: float = 1e6, + session: session.Session = session.global_session, + ): + self._reader_cache_id = "" + self.NAMESPACE: Tuple[str, str] = ("", server) + super().__init__( + tstart, + tend, + max_freq, + session, + pulse=pulse, + server=server, + default_error=default_error, + ) + self.pulse = pulse + + def _get_charge_exchange( + self, + uid: str, + instrument: str, + revision: RevisionLike, + quantities: Set[str], + ) -> Dict[str, Any]: + + Rmin, Rmax = self.MACHINE_DIMS[0][0], self.MACHINE_DIMS[0][1] + zmin, zmax = self.MACHINE_DIMS[1][0], self.MACHINE_DIMS[1][1] + + results: Dict[str, Any] = { + "length": np.random.randint(4, 20), + "machine_dims": self.MACHINE_DIMS, + } + dt = np.random.uniform(0.001, 1.0) + times = np.arange(TSTART, TEND, dt) + wavelength = np.arange(520, 530, 0.1) + nt = times.shape[0] + results["times"] = times + results["wavelength"] = wavelength + results["spectra"] = np.random.uniform( + 10, 10.0e3, (nt, results["length"], wavelength.size) + ) + results["fit"] = deepcopy(results["spectra"]) + results["texp"] = np.full_like(times, dt) + + results["location"] = np.array([[1.0, 2.0, 3.0]] * results["length"]) + results["direction"] = np.array([[1.0, 2.0, 3.0]] * results["length"]) + + results["element"] = "element" + results["revision"] = np.random.randint(0, 10) + results["x"] = np.random.uniform(Rmin, Rmax, (results["length"],)) + results["y"] = np.random.uniform(Rmin, Rmax, (results["length"],)) + results["z"] = np.random.uniform(zmin, zmax, (results["length"],)) + results["R"] = np.random.uniform(Rmin, Rmax, (results["length"],)) + results["ti"] = np.random.uniform(10, 10.0e3, (nt, results["length"])) + results["vtor"] = np.random.uniform(1.0e2, 1.0e6, (nt, results["length"])) + results["angf"] = np.random.uniform(1.0e2, 1.0e6, (nt, results["length"])) + results["conc"] = np.random.uniform(1.0e-6, 1.0e-1, (nt, results["length"])) + results["bad_channels"] = [] + + for quantity in quantities: + results[f"{quantity}_records"] = [ + f"{quantity}_R_path", + f"{quantity}_z_path", + f"{quantity}_element_path", + f"{quantity}_time_path", + f"{quantity}_ti_path", + f"{quantity}_angf_path", + f"{quantity}_conc_path", + ] + + return results + + def _get_spectrometer( + self, + uid: str, + instrument: str, + revision: RevisionLike, + quantities: Set[str], + ) -> Dict[str, Any]: + + results: Dict[str, Any] = { + "length": np.random.randint(4, 20), + "machine_dims": self.MACHINE_DIMS, + } + dt = np.random.uniform(0.001, 1.0) + times = np.arange(TSTART, TEND, dt) + wavelength = np.arange(520, 530, 0.1) + nt = times.shape[0] + results["times"] = times + results["wavelength"] = wavelength + results["spectra"] = np.random.uniform( + 10, 10.0e3, (nt, results["length"], wavelength.size) + ) + + results["location"] = np.array([[1.0, 2.0, 3.0]] * results["length"]) + results["direction"] = np.array([[1.0, 2.0, 3.0]] * results["length"]) + + results["revision"] = np.random.randint(0, 10) + + for quantity in quantities: + results[f"{quantity}_records"] = [ + f"{quantity}_time_path", + f"{quantity}_spectra_path", + ] + + return results + + def _get_thomson_scattering( + self, + uid: str, + instrument: str, + revision: RevisionLike, + quantities: Set[str], + ) -> Dict[str, Any]: + + Rmin, Rmax = self.MACHINE_DIMS[0][0], self.MACHINE_DIMS[0][1] + zmin, zmax = self.MACHINE_DIMS[1][0], self.MACHINE_DIMS[1][1] + + results: Dict[str, Any] = { + "length": np.random.randint(4, 20), + "machine_dims": self.MACHINE_DIMS, + } + + dt = np.random.uniform(0.001, 1.0) + times = np.arange(TSTART, TEND, dt) + nt = times.shape[0] + results["times"] = times + results["revision"] = np.random.randint(0, 10) + + results["x"] = np.random.uniform(Rmin, Rmax, (results["length"],)) + results["y"] = np.random.uniform(Rmin, Rmax, (results["length"],)) + results["z"] = np.random.uniform(zmin, zmax, (results["length"],)) + results["R"] = np.random.uniform(Rmin, Rmax, (results["length"],)) + results["chi2"] = np.random.uniform(0, 2.0, (nt, results["length"])) + results["te"] = np.random.uniform(10, 10.0e3, (nt, results["length"])) + results["ne"] = np.random.uniform(1.0e16, 1.0e21, (nt, results["length"])) + results["bad_channels"] = [] + + for quantity in quantities: + results[f"{quantity}_records"] = [ + f"{quantity}_Rz_path", + f"{quantity}_value_path", + f"{quantity}_error_path", + ] + + return results + + def _get_equilibrium( + self, + uid: str, + calculation: str, + revision: RevisionLike, + quantities: Set[str], + ) -> Dict[str, Any]: + + Rmin, Rmax = self.MACHINE_DIMS[0][0], self.MACHINE_DIMS[0][1] + zmin, zmax = self.MACHINE_DIMS[1][0], self.MACHINE_DIMS[1][1] + + results: Dict[str, Any] = {} + + dt = np.random.uniform(0.001, 1.0) + times = np.arange(TSTART, TEND, dt) + results["times"] = times + nt = times.shape[0] + nrho = np.random.randint(20, 40) + + results["element"] = "element" + + results["R"] = np.random.uniform(Rmin, Rmax, (nrho,)) + results["z"] = np.random.uniform(zmin, zmax, (nrho,)) + + results["rgeo"] = np.random.uniform(Rmin, Rmax, (nt,)) + results["rmag"] = np.random.uniform(Rmin, Rmax, (nt,)) + results["zmag"] = np.random.uniform(zmin, zmax, (nt,)) + results["ipla"] = np.random.uniform(1.0e4, 1.0e6, (nt,)) + results["wp"] = np.random.uniform(1.0e3, 1.0e5, (nt,)) + results["df"] = np.random.uniform(0, 1, (nt,)) + results["faxs"] = np.random.uniform(1.0e-6, 0.1, (nt,)) + results["fbnd"] = np.random.uniform(-1, 1, (nt,)) + + results["psin"] = np.random.uniform(0, 1, (nrho,)) + results["psi_r"] = np.random.uniform(Rmin, Rmax, (nrho,)) + results["psi_z"] = np.random.uniform(zmin, zmax, (nrho,)) + + results["f"] = np.random.uniform(1.0e-6, 0.1, (nt, nrho)) + results["ftor"] = np.random.uniform(1.0e-4, 1.0e-2, (nt, nrho)) + results["ajac"] = np.random.uniform(1.0e-3, 2.0, (nt, nrho)) + results["vjac"] = np.random.uniform(1.0e-3, 2.0, (nt, nrho)) + results["rmji"] = np.random.uniform(Rmin, Rmax, (nt, nrho)) + results["rmjo"] = np.random.uniform(Rmin, Rmax, (nt, nrho)) + results["rbnd"] = np.random.uniform(Rmin, Rmax, (nt, nrho)) + results["zbnd"] = np.random.uniform(zmin, zmax, (nt, nrho)) + results["psi"] = np.random.uniform(-1, 1, (nt, nrho, nrho)) + + results["revision"] = np.random.randint(0, 10) + + for quantity in quantities: + results[f"{quantity}_records"] = [f"{quantity}_records"] + results["psi_records"] = ["value_records", "R_records", "z_records"] + return results + + def _get_radiation( + self, + uid: str, + instrument: str, + revision: RevisionLike, + quantities: Set[str], + ) -> Dict[str, Any]: + + results: Dict[str, Any] = { + "length": np.random.randint(4, 20), + "machine_dims": self.MACHINE_DIMS, + } + + results["revision"] = np.random.randint(0, 10) + dt = np.random.uniform(0.001, 1.0) + times = np.arange(TSTART, TEND, dt) + results["times"] = times + nt = times.shape[0] + results["location"] = np.array([[1.0, 2.0, 3.0]] * results["length"]) + results["direction"] = np.array([[1.0, 2.0, 3.0]] * results["length"]) + + results["times"] = times + results["brightness"] = np.random.uniform(0, 1.0e6, (nt, results["length"])) + + results["brightness_records"] = ["brightness_records"] * results["length"] + + return results + + def _get_bremsstrahlung_spectroscopy( + self, + uid: str, + instrument: str, + revision: RevisionLike, + quantities: Set[str], + ) -> Dict[str, Any]: + + results: Dict[str, Any] = { + "length": np.random.randint(4, 20), + "machine_dims": self.MACHINE_DIMS, + } + results["revision"] = np.random.randint(0, 10) + dt = np.random.uniform(0.001, 1.0) + times = np.arange(TSTART, TEND, dt) + results["times"] = times + nt = times.shape[0] + + results["location"] = np.array([[1.0, 2.0, 3.0]] * results["length"]) + results["direction"] = np.array([[1.0, 2.0, 3.0]] * results["length"]) + + quantity = "zeff" + results[quantity] = np.random.uniform(0, 1.0e6, (nt, results["length"])) + + results[f"{quantity}_records"] = [ + f"{quantity}_path_records", + f"{quantity}_los_records", + ] + + return results + + def _get_helike_spectroscopy( + self, + uid: str, + instrument: str, + revision: RevisionLike, + quantities: Set[str], + ) -> Dict[str, Any]: + + nwavelength = np.random.randint(256, 1024) + wavelength_start, wavelength_end = 3.8, 4.0 + + results: Dict[str, Any] = { + "length": np.random.randint(4, 20), + "machine_dims": self.MACHINE_DIMS, + } + results["revision"] = np.random.randint(0, 10) + dt = np.random.uniform(0.001, 1.0) + times = np.arange(TSTART, TEND, dt) + results["times"] = times + nt = times.shape[0] + + results["location"] = np.array([[1.0, 2.0, 3.0]] * results["length"]) + results["direction"] = np.array([[1.0, 2.0, 3.0]] * results["length"]) + + results["wavelength"] = np.linspace( + wavelength_start, wavelength_end, nwavelength + ) + + for quantity in quantities: + if quantity == "spectra": + results[quantity] = np.random.uniform( + 0, 1.0e6, (nt, results["length"], nwavelength) + ) + else: + results[quantity] = np.random.uniform(0, 1.0e4, (nt, results["length"])) + + results[f"{quantity}_records"] = [ + f"{quantity}_path_records", + ] + + return results + + def _get_diode_filters( + self, + uid: str, + instrument: str, + revision: RevisionLike, + quantities: Set[str], + ) -> Dict[str, Any]: + + results: Dict[str, Any] = { + "length": np.random.randint(4, 20), + "machine_dims": self.MACHINE_DIMS, + } + + results["revision"] = np.random.randint(0, 10) + dt = np.random.uniform(0.001, 1.0) + times = np.arange(TSTART, TEND, dt) + results["times"] = times + nt = times.shape[0] + + results["location"] = np.array([[1.0, 2.0, 3.0]] * results["length"]) + results["direction"] = np.array([[1.0, 2.0, 3.0]] * results["length"]) + results["labels"] = np.array(["label"] * results["length"]) + + for quantity in quantities: + results[quantity] = np.random.uniform(0, 1.0e6, (nt, results["length"])) + + results[f"{quantity}_records"] = [ + f"{quantity}_path_records", + ] + results[f"{quantity}_error_records"] = [ + f"{quantity}_error_path_records", + ] + + return results + + def _get_interferometry( + self, + uid: str, + instrument: str, + revision: RevisionLike, + quantities: Set[str], + ) -> Dict[str, Any]: + + results: Dict[str, Any] = { + "length": np.random.randint(4, 20), + "machine_dims": self.MACHINE_DIMS, + } + + results["revision"] = np.random.randint(0, 10) + dt = np.random.uniform(0.001, 1.0) + times = np.arange(TSTART, TEND, dt) + results["times"] = times + nt = times.shape[0] + + results["location"] = np.array([[1.0, 2.0, 3.0]] * results["length"]) + results["direction"] = np.array([[1.0, 2.0, 3.0]] * results["length"]) + for quantity in quantities: + results[quantity] = np.random.uniform(0, 1.0e6, (nt, results["length"])) + + results[f"{quantity}_records"] = [ + f"{quantity}_path_records", + ] + results[f"{quantity}_error_records"] = [ + f"{quantity}_error_path_records", + ] + + return results + + def _get_bad_channels( + self, uid: str, instrument: str, quantity: str + ) -> List[Number]: + return [] + + def _set_times_item( + self, + results: Dict[str, Any], + times: np.ndarray, + ): + if "times" not in results: + times = times + + def close(self): + del self._client + + def requires_authentication(self): + return True + + +def _test_get_methods( + instrument="ts", + nsamples=1, +): + """ + Generalised test for all get methods of the abstractreader + """ + + for i in range(nsamples): + reader = Reader( + 1, + TSTART, + TEND, + ) + + quantities = set(AVAILABLE_QUANTITIES[reader.INSTRUMENT_METHODS[instrument]]) + + results = reader.get("", instrument, 0, quantities) + + # Check whether data is as expected + for q, actual, expected in [(q, results[q], results[q]) for q in quantities]: + assert np.all(actual.values == expected) + + +def test_get_thomson_scattering(): + _test_get_methods(instrument="thomson_scattering", nsamples=10) + + +def test_get_charge_exchange(): + _test_get_methods(instrument="charge_exchange", nsamples=10) + + +def test_get_spectrometer(): + _test_get_methods(instrument="spectrometer", nsamples=10) + + +def test_get_equilibrium(): + _test_get_methods(instrument="equilibrium", nsamples=10) + + +def test_get_radiation(): + _test_get_methods(instrument="radiation", nsamples=10) + + +def test_get_bremsstrahlung_spectroscopy(): + _test_get_methods( + instrument="bremsstrahlung_spectroscopy", + nsamples=10, + ) + + +def test_get_helike_spectroscopy(): + _test_get_methods( + instrument="helike_spectroscopy", + nsamples=10, + ) + + +def test_get_diode_filters(): + _test_get_methods( + instrument="diode_filters", + nsamples=10, + ) + + +def test_get_interferometry(): + _test_get_methods( + instrument="interferometry", + nsamples=10, + ) From 52fd84bb3656fac0450323fda36af23bc05a540e Mon Sep 17 00:00:00 2001 From: Marco Sertoli Date: Fri, 27 Oct 2023 15:32:32 +0100 Subject: [PATCH 06/17] Deleting all unused stuff in st40 branch --- indica/operators/__init__.py | 6 - indica/operators/centrifugal_asymmetry.py | 274 ---------- indica/operators/main_ion_density.py | 105 ---- indica/operators/mean_charge.py | 101 ---- snippets/plot_jet.py | 20 - snippets/read_KT7D.py | 120 ----- snippets/test_get_efit.py | 26 - snippets/wall_coords_jet.txt | 251 --------- tests/unit/fake_equilibrium.py | 251 --------- .../readers/test_abstract_reader_pytest.py | 507 ------------------ tests/unit/test_offset.py | 0 tests/unit/test_session.py | 66 --- 12 files changed, 1727 deletions(-) delete mode 100644 indica/operators/centrifugal_asymmetry.py delete mode 100644 indica/operators/main_ion_density.py delete mode 100644 indica/operators/mean_charge.py delete mode 100644 snippets/plot_jet.py delete mode 100644 snippets/read_KT7D.py delete mode 100644 snippets/test_get_efit.py delete mode 100644 snippets/wall_coords_jet.txt delete mode 100644 tests/unit/fake_equilibrium.py delete mode 100644 tests/unit/readers/test_abstract_reader_pytest.py delete mode 100644 tests/unit/test_offset.py diff --git a/indica/operators/__init__.py b/indica/operators/__init__.py index 4bd71ea9..2ff71550 100644 --- a/indica/operators/__init__.py +++ b/indica/operators/__init__.py @@ -2,9 +2,6 @@ from .abstractoperator import OperatorError from .atomic_data import FractionalAbundance from .atomic_data import PowerLoss -from .centrifugal_asymmetry import AsymmetryParameter -from .centrifugal_asymmetry import ToroidalRotation -from .mean_charge import MeanCharge from .zeff import CalcZeff __all__ = [ @@ -12,8 +9,5 @@ "OperatorError", "FractionalAbundance", "PowerLoss", - "ToroidalRotation", - "AsymmetryParameter", - "MeanCharge", "CalcZeff", ] diff --git a/indica/operators/centrifugal_asymmetry.py b/indica/operators/centrifugal_asymmetry.py deleted file mode 100644 index 68d423b0..00000000 --- a/indica/operators/centrifugal_asymmetry.py +++ /dev/null @@ -1,274 +0,0 @@ -from typing import Any -from typing import List -from typing import Tuple -from typing import Union - -from xarray.core.dataarray import DataArray - -from .abstractoperator import EllipsisType -from .abstractoperator import Operator -from .. import session -from ..datatypes import DataType -from ..datatypes import ELEMENTS -from ..utilities import input_check - - -class ToroidalRotation(Operator): - """Calculate the toroidal rotation from asymmetry parameter. - - Parameters - ---------- - - Returns - ------- - toroidal_rotation - xarray.DataArray containing toroidal rotation for a given impurity element - - Attributes - ---------- - ARGUMENT_TYPES: List[DataType] - Ordered list of the types of data expected for each argument of the - operator. - RESULT_TYPES: List[DataType] - Ordered list of the types of data returned by the operator. - """ - - ARGUMENT_TYPES: List[Union[DataType, EllipsisType]] = [] - - RESULT_TYPES: List[Union[DataType, EllipsisType]] = [ - ("toroidal_rotation", "plasma"), - ] - - def __init__(self, sess: session.Session = session.global_session): - super().__init__(sess=sess) - - def return_types(self, *args: DataType) -> Tuple[Any, ...]: - return super().return_types(*args) - - def __call__( # type: ignore - self, - asymmetry_parameters: DataArray, - ion_temperature: DataArray, - main_ion: str, - impurity: str, - Zeff: DataArray, - electron_temp: DataArray, - ): - """Calculates the toroidal rotation frequency from the asymmetry parameter. - - Parameters - ---------- - asymmetry_parameters - xarray.DataArray containing asymmetry parameters data. In units of m^-2. - ion_temperature - xarray.DataArray containing ion temperature data. In units of eV. - main_ion - Element symbol of main ion. - impurity - Element symbol of chosen impurity element. - Zeff - xarray.DataArray containing Z-effective data from diagnostics. - electron_temp - xarray.DataArray containing electron temperature data. In units of eV. - - Returns - ------- - toroidal_rotation - xarray.DataArray containing data for toroidal rotation frequencies - for the given impurity element - """ - input_check( - "asymmetry_parameters", - asymmetry_parameters, - DataArray, - ndim_to_check=3, - positive=False, - strictly_positive=False, - ) - - input_check( - "ion_temperature", - ion_temperature, - DataArray, - ndim_to_check=3, - strictly_positive=True, - ) - - input_check("main_ion", main_ion, str) - - try: - assert main_ion in list(ELEMENTS.keys()) - except AssertionError: - raise ValueError(f"main_ion must be one of {list(ELEMENTS.keys())}") - - input_check("impurity", impurity, str) - - try: - assert impurity in list(ELEMENTS.keys()) - except AssertionError: - raise ValueError(f"impurity must be one of {list(ELEMENTS.keys())}") - - input_check("Zeff", Zeff, DataArray, ndim_to_check=2, strictly_positive=False) - - input_check( - "electron_temp", - electron_temp, - DataArray, - ndim_to_check=2, - strictly_positive=True, - ) - - asymmetry_parameter = asymmetry_parameters.sel(elements=impurity) - - impurity_mass_int = ELEMENTS[impurity][1] - - unified_atomic_mass_unit = 931.4941e6 # in eV/c^2 - impurity_mass = float(impurity_mass_int) * unified_atomic_mass_unit - - mean_charge = ELEMENTS[impurity][0] - - main_ion_mass_int = ELEMENTS[main_ion][1] - - main_ion_mass = float(main_ion_mass_int) * unified_atomic_mass_unit - - ion_temperature = ion_temperature.sel(elements=impurity) - - # mypy on the github CI suggests that * is an Unsupported operand type - # between float and DataArray, don't know how to fix yet so for now ignored - toroidal_rotation = 2.0 * ion_temperature * asymmetry_parameter # type: ignore - toroidal_rotation /= impurity_mass * ( - 1.0 - - (mean_charge * main_ion_mass * Zeff * electron_temp) # type: ignore - / (impurity_mass * (ion_temperature + Zeff * electron_temp)) - ) - - toroidal_rotation = toroidal_rotation**0.5 - - c = 3.0e8 # speed of light in vacuum - toroidal_rotation *= c - - return toroidal_rotation - - -class AsymmetryParameter(Operator): - """Calculate the asymmetry parameter from toroidal rotation. - - Parameters - ---------- - - Returns - ------- - asymmetry_parameter - xarray.DataArray containing asymmetry_parameter for a given impurity element - - """ - - ARGUMENT_TYPES: List[Union[DataType, EllipsisType]] = [] - - def __init__(self, sess: session.Session = session.global_session): - super().__init__(sess=sess) - - def return_types(self, *args: DataType) -> Tuple[Any, ...]: - return super().return_types(*args) - - def __call__( # type: ignore - self, - toroidal_rotations: DataArray, - ion_temperature: DataArray, - main_ion: str, - impurity: str, - Zeff: DataArray, - electron_temp: DataArray, - ): - """Calculates the asymmetry parameter from the toroidal rotation frequency. - - Parameters - ---------- - toroidal_rotations - xarray.DataArray containing toroidal rotation frequencies data. - In units of ms^-1. - ion_temperature - xarray.DataArray containing ion temperature data. In units of eV. - main_ion - Element symbol of main ion. - impurity - Element symbol of chosen impurity element. - Zeff - xarray.DataArray containing Z-effective data from diagnostics. - electron_temp - xarray.DataArray containing electron temperature data. In units of eV. - - Returns - ------- - asymmetry_parameter - xarray.DataArray containing data for asymmetry parameters - for the given impurity element - """ - input_check( - "toroidal_rotations", - toroidal_rotations, - DataArray, - ndim_to_check=3, - strictly_positive=False, - ) - - input_check( - "ion_temperature", - ion_temperature, - DataArray, - ndim_to_check=3, - strictly_positive=True, - ) - - input_check("main_ion", main_ion, str) - - try: - assert main_ion in list(ELEMENTS.keys()) - except AssertionError: - raise ValueError(f"main_ion must be one of {list(ELEMENTS.keys())}") - - input_check("impurity", impurity, str) - - try: - assert impurity in list(ELEMENTS.keys()) - except AssertionError: - raise ValueError(f"impurity must be one of {list(ELEMENTS.keys())}") - - input_check("Zeff", Zeff, DataArray, ndim_to_check=2, strictly_positive=False) - - input_check( - "electron_temp", - electron_temp, - DataArray, - ndim_to_check=2, - strictly_positive=True, - ) - - toroidal_rotations = toroidal_rotations.sel(elements=impurity) - - impurity_mass_int = ELEMENTS[impurity][1] - - unified_atomic_mass_unit = 931.4941e6 # in eV/c^2 - impurity_mass = float(impurity_mass_int) * unified_atomic_mass_unit - - mean_charge = ELEMENTS[impurity][0] - - main_ion_mass_int = ELEMENTS[main_ion][1] - - main_ion_mass = float(main_ion_mass_int) * unified_atomic_mass_unit - - ion_temperature = ion_temperature.sel(elements=impurity) - - c = 3.0e8 # speed of light in m/s - toroidal_rotations /= c - - # mypy on the github CI suggests that * is in an Unsupported operand type - # between float and DataArray, don't know how to fix yet so for now ignored - asymmetry_parameter = ( - impurity_mass * (toroidal_rotations**2) / (2.0 * ion_temperature) - ) - asymmetry_parameter *= 1.0 - ( - mean_charge * main_ion_mass * Zeff * electron_temp # type: ignore - ) / (impurity_mass * (ion_temperature + Zeff * electron_temp)) - - return asymmetry_parameter diff --git a/indica/operators/main_ion_density.py b/indica/operators/main_ion_density.py deleted file mode 100644 index 2d27b5bc..00000000 --- a/indica/operators/main_ion_density.py +++ /dev/null @@ -1,105 +0,0 @@ -"""Operator calculating the main ion density given the densities of impurities. -""" - -from typing import List -from typing import Tuple -from typing import Union - -from xarray.core.dataarray import DataArray - -from indica.datatypes import DataType -from .abstractoperator import EllipsisType -from .abstractoperator import Operator -from .. import session -from ..utilities import input_check - - -class MainIonDensity(Operator): - """Calculates the main ion density from given impurity densities and mean charge. - - Attributes - ---------- - ARGUMENT_TYPES: List[DataType] - Ordered list of the types of data expected for each argument of the - operator. - RESULT_TYPES: List[DataType] - Ordered list of the types of data returned by the operator. - - Returns - ------- - main_ion_density - xarray.DataArray of the main ion density. - - Methods - ------- - __call__(impurity_densities, electron_density, mean_charge) - Calculates the main ion density from given impurity densities and mean charge. - """ - - ARGUMENT_TYPES: List[Union[DataType, EllipsisType]] = [] - - RESULT_TYPES: List[Union[DataType, EllipsisType]] = [ - ("main_ion", "number_density"), - ] - - def __init__(self, sess: session.Session = session.global_session): - super().__init__(sess=sess) - - def return_types(self, *args: DataType) -> Tuple[DataType, ...]: - return super().return_types(*args) - - def __call__( # type: ignore - self, - impurity_densities: DataArray, - electron_density: DataArray, - mean_charge: DataArray, - ): - """Calculates the main ion density from given impurity densities - and mean charge. - - Parameters - ---------- - impurity_densities - xarray.DataArray of impurity densities for all impurity elements - of interest. - electron_density - xarray.DataArray of electron density - mean_charge - xarray.DataArray of mean charge of all impurity elements of interest. - This can be provided manually - (with dimensions of ["element", "rho_poloidal", "t]), - or can be passed as the results of MeanCharge.__call__ - - Returns - ------- - main_ion_density - xarray.DataArray of the main ion density. - """ - # no ndim check since impurity densities can have coords: - # [elements, rho, t] or [elements, R, z, t] - input_check( - "impurity_densities", - impurity_densities, - DataArray, - strictly_positive=False, - ) - - input_check( - "electron_density", - electron_density, - DataArray, - strictly_positive=False, - ) - - input_check( - "mean_charge", - mean_charge, - DataArray, - strictly_positive=False, - ) - - main_ion_density = electron_density - (mean_charge * impurity_densities).sum( - "element" - ) - - return main_ion_density diff --git a/indica/operators/mean_charge.py b/indica/operators/mean_charge.py deleted file mode 100644 index 8cd02a1d..00000000 --- a/indica/operators/mean_charge.py +++ /dev/null @@ -1,101 +0,0 @@ -"""Operator calculating the mean charge from the fractional abundance of all -ionisation charges of a given element. -""" - -from typing import List -from typing import Tuple -from typing import Union - -import numpy as np -from xarray.core.common import zeros_like -from xarray.core.dataarray import DataArray - -from .abstractoperator import EllipsisType -from .abstractoperator import Operator -from .. import session -from ..datatypes import DataType -from ..datatypes import ELEMENTS -from ..utilities import input_check - - -class MeanCharge(Operator): - """Calculate mean charge for a given element from its fractional abundance. - - Parameters - ---------- - - Returns - ------- - mean_charge - numpy.ndarray of mean charge of the given element. - - """ - - ARGUMENT_TYPES: List[Union[DataType, EllipsisType]] = [] - - RESULT_TYPES: List[Union[DataType, EllipsisType]] = [ - ("mean_charge", "impurity_element"), - ] - - def __init__(self, sess: session.Session = session.global_session): - super().__init__(sess=sess) - - def return_types(self, *args: DataType) -> Tuple[DataType, ...]: - return super().return_types(*args) - - def __call__(self, FracAbundObj: DataArray, element: str): # type: ignore - """Function to calculate the mean charge. - - Parameters - ---------- - FracAbundObj - numpy.ndarray describing the fractional abundance of the given element. - The first axis must correspond to the ionisation charges of the element. - element - Symbol of the element for which the mean charge is desired. - - Returns - ------- - mean_charge - numpy.ndarray of mean charge of the given element. - """ - - input_check( - "FracAbundObj", - FracAbundObj, - DataArray, - ndim_to_check=3, - strictly_positive=False, - ) - input_check("element", element, str) - - try: - assert element in ELEMENTS.keys() - except AssertionError: - raise ValueError( - f"Please input a single valid element from list:\ - {list(ELEMENTS.keys())}" - ) - - element_atomic_number = ELEMENTS[element][0] - - ionisation_charges = np.arange(0, element_atomic_number + 1) # type: ignore - - try: - assert ionisation_charges.shape[0] == FracAbundObj.shape[0] - except AssertionError: - raise AssertionError( - f"Number of ionisation charges in the \ - FractionalAbundance object do not match the expected number for the \ - element provided, {element}" - ) - - mean_charge = zeros_like(FracAbundObj) - mean_charge = mean_charge.isel(ion_charges=0) - mean_charge.drop_vars("ion_charges") - - mean_charge = np.sum( - ionisation_charges[:, np.newaxis, np.newaxis] * FracAbundObj, axis=0 - ) - - return mean_charge diff --git a/snippets/plot_jet.py b/snippets/plot_jet.py deleted file mode 100644 index 53248f95..00000000 --- a/snippets/plot_jet.py +++ /dev/null @@ -1,20 +0,0 @@ -import matplotlib.pylab as plt -import numpy as np - - -def plot_jet_wall( - filewall="/home/msertoli/work/datapool/wall_coords_jet.txt", divertor=True -): - """ - Plot JET first wall poloidal cross-section - """ - data = np.genfromtxt(filewall) - - plt.figure() - plt.plot(data[:, 0], data[:, 1], color="black", linewidth=2.5) - if divertor: - plt.xlim(2.3, 3.0) - plt.ylim(-1.8, -1.3) - plt.gca().set_aspect("equal", adjustable="box") - plt.xlabel("R (m)") - plt.ylabel("z (m)") diff --git a/snippets/read_KT7D.py b/snippets/read_KT7D.py deleted file mode 100644 index 6b63bd55..00000000 --- a/snippets/read_KT7D.py +++ /dev/null @@ -1,120 +0,0 @@ -import getdat as gd -import numpy as np - - -def get_kt7c(pulse, skip=None, trim40=False, backgrd=True, foregrd=False): - - if skip is None: - skip = 5 - - # get data - nwds = 0 - data_node = "df/t7-spec:003" - data, tvec, nwds, title, units, ier = gd.getdat(data_node, pulse, nwds=nwds) - ndata = nwds - - # get time vector - nwds_tim = 0 - tvec_tim, nwds_tim, ier_tim = gd.gettim(data_node, pulse, nwds=nwds_tim) - - # get number of frames - nframe_node = "DF/T7-NFRAME npixel * nframe: - data = data[0 : npixel * nframe] - elif ndata < npixel * nframe: - nframe = ndata // npixel - data = data[0 : npixel * nframe] - - # calculate exposure time - treadout = tvec_tim[0:nframe] - exp_time = treadout * 0.0 - exp_time[1:] = treadout[1:] - treadout[0:-1] - exp_time[0] = exp_time[1] # bald-faced lie (unknown exp for first frame) - exp_time = exp_time.round(decimals=4) - time = treadout - (exp_time / 2.0) - - # reshape data array - shape = (nframe, npixel) - data = np.reshape(data, shape) - - # keep only used pixels - npixel = 1024 - pixel = np.arange(npixel) - data = data[:, 0:npixel] - - # Calculate background or foreground: - # - background: exclude first 5 frames & the 2 'move' frames just before 40s - # - foreground: average last 20 frames - # In both cases normalise to exp_time (can vary during the discharge) - bkd = np.zeros(npixel) - if backgrd and not foregrd: - ind_bkd = np.argwhere(time < 40.0)[5:-2].flatten() - bkd = [] - for i in ind_bkd: - bkd.append(data[i, :] / exp_time[i]) - bkd = np.array(bkd).mean(axis=0).flatten() - - fgd = np.zeros(npixel) - if foregrd: - ind_fgd = range(nframe - 20, nframe) - fgd = [] - for i in ind_fgd: - fgd.append(data[i, :] / exp_time[i]) - fgd = np.array(fgd).mean(axis=0).flatten() - - # Normalise to exposure time (can vary during the discharge) and - # subtract background/foreground - for i in range(len(time)): - data[i, :] = data[i, :] / exp_time[i] - bkd - fgd - - # Retain data only for time > 40 s - if trim40: - indx40 = np.argwhere(time > 40.0).flatten() - data = data[indx40, :] - time = time[indx40] - treadout = treadout[indx40] - exp_time = exp_time[indx40] - - # Moc-up wavelength calibration: pixel offset accounts for changes - # in detector position - if pulse >= 80685: - pix_offs = +149 - elif pulse >= 80724: - pix_offs = 0 - elif pulse >= 84931: - pix_offs = -24 - elif pulse >= 89473: - pix_offs = -2 - else: - pix_offs = 0 - - c = [0.0] * 3 - c[0] = 4.0351381 - c[1] = 0.0033944632 - c[2] = -3.4947697e-07 - - # wavelength (nm) - pix = pixel + pix_offs - wave = c[0] + c[1] * pix + c[2] * pix**2 - - out = { - "data": data, - "time": time, - "treadout": treadout, - "exptime": exp_time, - "pixel": pixel, - "wave": wave, - } - - return out diff --git a/snippets/test_get_efit.py b/snippets/test_get_efit.py deleted file mode 100644 index 1a304299..00000000 --- a/snippets/test_get_efit.py +++ /dev/null @@ -1,26 +0,0 @@ -import getpass - -from indica.readers import PPFReader - -reader = PPFReader(90279, 45.0, 50.0) -if reader.requires_authentication: - user = input("JET username: ") - password = getpass.getpass("JET password: ") - assert reader.authenticate(user, password) - -data = reader.get("jetppf", "efit", 0) -print("f value") -print("=======\n") -print(data["f"]) -print("\n\n") -print("psi") -print("===\n") -print(data["psi"]) -print("\n\n") -print("Rmag") -print("====\n") -print(data["rmag"]) -print("\n\n") -print("Rbnd") -print("====\n") -print(data["rbnd"]) diff --git a/snippets/wall_coords_jet.txt b/snippets/wall_coords_jet.txt deleted file mode 100644 index b5b65582..00000000 --- a/snippets/wall_coords_jet.txt +++ /dev/null @@ -1,251 +0,0 @@ -3.28315 -1.12439 -3.31192 -1.08316 -3.32836 -1.06314 -3.35244 -1.03867 -3.37323 -1.0172 -3.42329 -0.95992 -3.44728 -0.93229 -3.49433 -0.87186 -3.5168 -0.84284 -3.55911 -0.78155 -3.58053 -0.75015 -3.61927 -0.68679 -3.63903 -0.6539 -3.67364 -0.58769 -3.68796 -0.56088 -3.7223 -0.48496 -3.73497 -0.4576 -3.76821 -0.37137 -3.77584 -0.35239 -3.80258 -0.2692 -3.81082 -0.24458 -3.83331 -0.15797 -3.83964 -0.13481 -3.85652 -0.05075 -3.86207 -0.02427 -3.87519 0.06852 -3.87835 0.08885 -3.88567 0.17625 -3.88793 0.1963 -3.89046 0.29045 -3.89144 0.31339 -3.88863 0.40756 -3.8884 0.42674 -3.88031 0.51993 -3.87894 0.53958 -3.86556 0.63228 -3.86349 0.64971 -3.84449 0.74347 -3.84091 0.76273 -3.81715 0.85325 -3.81303 0.87053 -3.78365 0.96125 -3.77904 0.97699 -3.74408 1.06719 -3.73821 1.08341 -3.69868 1.17046 -3.69176 1.18659 -3.67489 1.2188 -3.66741 1.23638 -3.6373 1.33388 -3.64211 1.40768 -3.62136 1.42656 -3.62076 1.4271 -3.59407 1.45138 -3.59348 1.45192 -3.56679 1.4762 -3.5662 1.47674 -3.53951 1.50102 -3.53892 1.50156 -3.51223 1.52584 -3.51164 1.52638 -3.48495 1.55066 -3.48436 1.55119 -3.45767 1.57547 -3.45708 1.57601 -3.43039 1.60029 -3.4298 1.60083 -3.40311 1.62511 -3.40252 1.62565 -3.38176 1.64453 -3.33154 1.70412 -3.28182 1.73872 -3.18634 1.81753 -3.13665 1.85212 -3.00098 1.88341 -2.86885 1.93955 -2.85247 1.94527 -2.77562 1.96797 -2.75763 1.97196 -2.67896 1.98285 -2.65886 1.98403 -2.57787 1.9829 -2.56057 1.98112 -2.48011 1.96774 -2.46596 1.96418 -2.38844 1.93892 -2.37422 1.93306 -2.29805 1.89475 -2.28755 1.88838 -2.19539 1.82284 -2.18241 1.82372 -2.1657 1.79082 -2.16533 1.7901 -2.14899 1.75794 -2.14863 1.75722 -2.13229 1.72506 -2.13193 1.72434 -2.11558 1.69217 -2.11522 1.69146 -2.09888 1.65929 -2.09852 1.65858 -2.08218 1.62641 -2.08181 1.6257 -2.06816 1.59882 -2.06756 1.59819 -2.05467 1.5645 -2.05438 1.56375 -2.04148 1.53005 -2.0412 1.52931 -2.0283 1.49561 -2.02801 1.49486 -2.01512 1.46117 -2.01483 1.46042 -2.00193 1.42672 -2.00164 1.42597 -1.98875 1.39228 -1.98846 1.39153 -1.97556 1.35783 -1.97528 1.35709 -1.9613 1.32058 -1.92992 1.27302 -1.92697 1.26095 -1.92259 1.25401 -1.94246 1.23457 -1.92727 1.15832 -1.92496 1.13801 -1.91241 1.06067 -1.90969 1.03999 -1.89845 0.96207 -1.89568 0.94009 -1.88577 0.86238 -1.8830 0.84067 -1.87309 0.76296 -1.87038 0.74168 -1.86048 0.66408 -1.85836 0.64287 -1.84987 0.56549 -1.84816 0.54523 -1.84184 0.4671 -1.84095 0.44616 -1.83724 0.36833 -1.83704 0.34737 -1.83592 0.27002 -1.83642 0.2485 -1.83787 0.17099 -1.83907 0.14967 -1.84313 0.07136 -1.84502 0.05067 -1.85168 -0.02749 -1.85424 -0.04782 -1.86338 -0.12484 -1.86665 -0.14554 -1.87848 -0.22313 -1.88235 -0.24315 -1.89673 -0.32025 -1.90128 -0.34019 -1.91822 -0.41681 -1.92341 -0.43655 -1.94275 -0.51204 -1.94865 -0.53185 -1.97056 -0.60693 -1.95973 -0.62652 -1.96175 -0.65756 -2.00911 -0.78399 -2.02042 -0.8113 -2.02073 -0.81204 -2.03454 -0.84537 -2.03484 -0.84612 -2.04865 -0.87945 -2.04896 -0.88019 -2.06276 -0.91352 -2.06307 -0.91426 -2.07688 -0.9476 -2.07718 -0.94834 -2.09099 -0.98167 -2.0913 -0.98241 -2.10511 -1.01574 -2.10549 -1.01667 -2.1193 -1.0500 -2.1196 -1.05074 -2.13341 -1.08407 -2.13372 -1.08482 -2.14752 -1.11815 -2.14783 -1.11889 -2.16164 -1.15222 -2.16194 -1.15296 -2.17575 -1.1863 -2.17606 -1.18704 -2.18986 -1.22037 -2.19017 -1.22111 -2.20149 -1.24842 -2.14463 -1.27494 -2.29362 -1.31483 -2.29362 -1.33144 -2.29544 -1.33443 -2.35993 -1.33443 -2.39619 -1.37323 -2.40915 -1.4003 -2.41225 -1.42198 -2.41293 -1.43148 -2.41293 -1.46854 -2.41224 -1.47678 -2.40762 -1.50441 -2.39801 -1.51641 -2.41921 -1.59223 -2.42117 -1.61022 -2.4188 -1.64283 -2.41628 -1.6561 -2.40573 -1.68971 -2.31498 -1.7387 -2.35349 -1.7387 -2.37428 -1.73504 -2.42744 -1.71349 -2.44623 -1.70983 -2.52369 -1.70983 -2.52459 -1.69997 -2.55911 -1.65498 -2.55296 -1.63799 -2.57391 -1.6018 -2.63299 -1.61714 -2.63369 -1.61989 -2.6938 -1.6355 -2.69434 -1.63821 -2.75443 -1.65481 -2.75517 -1.65658 -2.81471 -1.67203 -2.81467 -1.70788 -2.80425 -1.71158 -2.85703 -1.71158 -2.87846 -1.71602 -2.93644 -1.74139 -2.95732 -1.74595 -2.98698 -1.74595 -2.89768 -1.68233 -2.88199 -1.62282 -2.88163 -1.5916 -2.90045 -1.51041 -2.89049 -1.49841 -2.88786 -1.48925 -2.88591 -1.47397 -2.88591 -1.43566 -2.88946 -1.41714 -2.90082 -1.39278 -2.91335 -1.37624 -2.96348 -1.33481 -3.00975 -1.33481 -3.06005 -1.29777 -3.19404 -1.21404 -3.20225 -1.20891 -3.30634 -1.20891 -3.28315 -1.12439 diff --git a/tests/unit/fake_equilibrium.py b/tests/unit/fake_equilibrium.py deleted file mode 100644 index 0ac6ea5d..00000000 --- a/tests/unit/fake_equilibrium.py +++ /dev/null @@ -1,251 +0,0 @@ -"""A subclass of :py:class:`indica.equilibrium.Equilibrium` which fakes -the implementation.""" - -from itertools import product -from unittest.mock import MagicMock - -from hypothesis.strategies import composite -from hypothesis.strategies import floats -from hypothesis.strategies import one_of -from hypothesis.strategies import sampled_from -import numpy as np -from xarray import DataArray - -from indica.equilibrium import Equilibrium - - -FLUX_TYPES = ["poloidal", "toroidal"] - - -@composite -def flux_types(draw): - return draw(sampled_from(FLUX_TYPES)) - - -class FakeEquilibrium(Equilibrium): - """A class which fakes the behaviour of an Equilibrium object. Flux - surface and magnetic fields are taken to vary in an elliptical profile - away from the magnetif axis. - - Fluxes have form $$r^2 = \\frac{(R-R_{mag})^2}{a^2} + - \\frac{(z-z_{mag})^2}{b^2},$$ where $r = \\rho^n(1 + \\alpha t)$ and $a$, $b$, - $n$ and $\\alpha$ are parameters specified by the user at instantiation. - - $B_{tot}$ varies according to a different equation: - $$B_{tot} = \\frac{(1 + \\alpha t) a}{1 + bR} + (z - z_{mag}).$$ - - Paramter values may be specified at instantiation using - keyword-arguments of the constructor. There are also default - values available for flux kinds ``poloidal``, ``toroidal``, and - the total magnetic field strength. - - Parameters - ---------- - Rmag : float - Major radius of the magnetic axis - zmag : float - Vertical position of the magnetic axis - kwargs : Dict[str, float] - Values for parameters describing the equilibrium profile. Keys take the - form ``_``. The ```` may be any - string which will be accpeted as a ``kind`` argument in methods such as - :py:meth:`flux_coords``, or ``Btot`` if the paremeter is describing the - profile of total magnetic field strength. The ```` may - be ``a``, ``b``, ``n``, or ``alpha``. - - """ - - DEFAULT_PARAMS = { - "poloidal_a": 0.5, - "poloidal_b": 1.0, - "poloidal_n": 1, - "poloidal_alpha": 0.01, - "toroidal_a": 0.7, - "toroidal_b": 1.4, - "toroidal_n": 1, - "toroidal_alpha": -0.00005, - "Btot_a": 1.0, - "Btot_b": 1.0, - "Btot_alpha": 0.001, - } - - def __init__( - self, - Rmag=3.0, - zmag=0.0, - default_t=DataArray([0.0, 5e3], dims="t"), - Bmax=1.0, - **kwargs - ): - ones = DataArray(np.ones_like(default_t), coords=[("t", default_t)]) - self.rmag = np.abs(Rmag) * ones - self.zmag = zmag * ones - self.parameters = kwargs - for k, v in self.DEFAULT_PARAMS.items(): - if k not in self.parameters: - self.parameters[k] = v - self.default_t = default_t - self.prov_id = MagicMock() - self.provenance = MagicMock() - self._session = MagicMock() - - def Btot(self, R, z, t=None): - if t is None: - t = self.default_t - zmag = self.zmag - else: - zmag = self.zmag.interp( - t=t, method="nearest", kwargs={"fill_value": "extrapolate"} - ) - return ( - (1 + self.parameters["Btot_alpha"] * t) - * self.parameters["Btot_a"] - / (1 + self.parameters["Btot_b"] * R) - + z - - zmag, - t, - ) - - def enclosed_volume(self, rho, t=None, kind="poloidal"): - if t is None: - t = self.default_t - rmag = self.rmag - else: - rmag = self.rmag.interp( - t=t, method="nearest", kwargs={"fill_value": "extrapolate"} - ) - a = self.parameters[kind + "_a"] - b = self.parameters[kind + "_b"] - n = self.parameters[kind + "_n"] - alpha = self.parameters[kind + "_alpha"] - vol = 2 * np.pi**2 * a * b * rho ** (2 * n) * (1 + alpha * t) ** 2 * rmag - - # blank area variable to return so that tests pass - area = 0 - return vol, t, area - - def invert_enclosed_volume(self, vol, t=None, kind="poloidal"): - if t is None: - t = self.default_t - rmag = self.rmag - else: - rmag = self.rmag.interp( - t=t, method="nearest", kwargs={"fill_value": "extrapolate"} - ) - a = self.parameters[kind + "_a"] - b = self.parameters[kind + "_b"] - n = self.parameters[kind + "_n"] - alpha = self.parameters[kind + "_alpha"] - rho = (vol / (2 * np.pi**2 * a * b * (1 + alpha * t) ** 2 * rmag)) ** ( - 0.5 / n - ) - return rho, t - - def R_hfs(self, rho, t=None, kind="poloidal"): - R, _, t = self.spatial_coords(rho, np.pi, t, kind) - return R, t - - def R_lfs(self, rho, t=None, kind="poloidal"): - R, _, t = self.spatial_coords(rho, 0.0, t, kind) - return R, t - - def minor_radius(self, rho, theta, t=None, kind="poloidal"): - if t is None: - t = self.default_t - r = rho ** self.parameters[kind + "_n"] * ( - 1 + self.parameters[kind + "_alpha"] * t - ) - return r, t - - def flux_coords(self, R, z, t=None, kind="poloidal"): - if t is None: - t = self.default_t - rmag = self.rmag - zmag = self.zmag - else: - rmag = self.rmag.interp( - t=t, method="nearest", kwargs={"fill_value": "extrapolate"} - ) - zmag = self.zmag.interp( - t=t, method="nearest", kwargs={"fill_value": "extrapolate"} - ) - rho = ( - ( - (R - rmag) ** 2 / self.parameters[kind + "_a"] ** 2 - + (z - zmag) ** 2 / self.parameters[kind + "_b"] ** 2 - ) - / (1 + self.parameters[kind + "_alpha"] * t) ** 2 - ) ** (1 / (2 * self.parameters[kind + "_n"])) - theta = np.arctan2((z - zmag), (R - rmag)) - return rho, theta, t - - def spatial_coords(self, rho, theta, t=None, kind="poloidal"): - if t is None: - t = self.default_t - rmag = self.rmag - zmag = self.zmag - else: - rmag = self.rmag.interp( - t=t, method="nearest", kwargs={"fill_value": "extrapolate"} - ) - zmag = self.zmag.interp( - t=t, method="nearest", kwargs={"fill_value": "extrapolate"} - ) - tan_theta = np.tan(theta) - dR = np.sign(np.cos(theta)) * np.sqrt( - ( - rho ** self.parameters[kind + "_n"] - * (1 + self.parameters[kind + "_alpha"] * t) - ) - ** 2 - / ( - 1 / self.parameters[kind + "_a"] ** 2 - + (tan_theta / self.parameters[kind + "_b"]) ** 2 - ) - ) - dz = tan_theta * dR - return rmag + dR, zmag + dz, t - - def convert_flux_coords( - self, rho, theta, t=None, from_kind="poloidal", to_kind="toroidal" - ): - R, z, t = self.spatial_coords(rho, theta, t, from_kind) - return self.flux_coords(R, z, t, to_kind) - - -@composite -def fake_equilibria( - draw, - Rmag, - zmag, - default_t=DataArray([0.0, 500.0], dims="t"), - flux_types=FLUX_TYPES, - **kwargs -): - """Generate instances of the FakeEquilibrium class, with the specified - flux types. Parameters will be drawn from the ``floats`` strategy, - unless explicitely specified as a keyword arguments. These - parameters should take the form ``_a``, - ``_b``, ``_n`` and ``_alpha``. In - addition to the flux types specified as an argument, you may - specify the parameter values for ``Btot``. - - """ - param_types = { - "a": floats(-0.9, 9.0).map(lambda x: x + 1.0), - "b": floats(-0.9, 9.0).map(lambda x: x + 1.0), - "n": one_of(sampled_from([1, 2, 0.5]), floats(0.2, 4.0)), - "alpha": floats(-0.02, 0.09).map(lambda x: x + 0.01), - } - param_values = kwargs - for flux, param in product(flux_types, param_types): - param_name = flux + "_" + param - if param_name not in param_values: - param_values[param_name] = draw(param_types[param]) - if "Btot_a" in flux_types and "Btot_a" not in param_values: - param_values["Btot_a"] = draw(floats(1e-3, 1e3)) - if "Btot_b" in flux_types and "Btot_b" not in param_values: - param_values["Btot_b"] = draw(floats(1e-5, 1e-2)) - if "Btot_alpha" in flux_types and "Btot_alpha" not in param_values: - param_values["Btot_alpha"] = draw(floats(-1e-3, 1e-3)) - return FakeEquilibrium(Rmag, zmag, default_t, **param_values) diff --git a/tests/unit/readers/test_abstract_reader_pytest.py b/tests/unit/readers/test_abstract_reader_pytest.py deleted file mode 100644 index cccd6aba..00000000 --- a/tests/unit/readers/test_abstract_reader_pytest.py +++ /dev/null @@ -1,507 +0,0 @@ -"""Test methods present on the base class DataReader.""" - -from copy import deepcopy -from numbers import Number -from typing import Any -from typing import Dict -from typing import List -from typing import Set -from typing import Tuple - -import numpy as np - -from indica import session -from indica.numpy_typing import RevisionLike -from indica.readers import DataReader -from indica.readers.available_quantities import AVAILABLE_QUANTITIES - - -# TODO these values should come from the machine dimensions variable of the reader - -TSTART = 0 -TEND = 10 - - -class Reader(DataReader): - """Class to read fake data""" - - MACHINE_DIMS = ((1.83, 3.9), (-1.75, 2.0)) - INSTRUMENT_METHODS = { - "thomson_scattering": "get_thomson_scattering", - "equilibrium": "get_equilibrium", - "cyclotron_emissions": "get_cyclotron_emissions", - "charge_exchange": "get_charge_exchange", - "spectrometer": "get_spectrometer", - "bremsstrahlung_spectroscopy": "get_bremsstrahlung_spectroscopy", - "radiation": "get_radiation", - "helike_spectroscopy": "get_helike_spectroscopy", - "interferometry": "get_interferometry", - "diode_filters": "get_diode_filters", - } - - def __init__( - self, - pulse: int, - tstart: float, - tend: float, - server: str = "", - default_error: float = 0.05, - max_freq: float = 1e6, - session: session.Session = session.global_session, - ): - self._reader_cache_id = "" - self.NAMESPACE: Tuple[str, str] = ("", server) - super().__init__( - tstart, - tend, - max_freq, - session, - pulse=pulse, - server=server, - default_error=default_error, - ) - self.pulse = pulse - - def _get_charge_exchange( - self, - uid: str, - instrument: str, - revision: RevisionLike, - quantities: Set[str], - ) -> Dict[str, Any]: - - Rmin, Rmax = self.MACHINE_DIMS[0][0], self.MACHINE_DIMS[0][1] - zmin, zmax = self.MACHINE_DIMS[1][0], self.MACHINE_DIMS[1][1] - - results: Dict[str, Any] = { - "length": np.random.randint(4, 20), - "machine_dims": self.MACHINE_DIMS, - } - dt = np.random.uniform(0.001, 1.0) - times = np.arange(TSTART, TEND, dt) - wavelength = np.arange(520, 530, 0.1) - nt = times.shape[0] - results["times"] = times - results["wavelength"] = wavelength - results["spectra"] = np.random.uniform( - 10, 10.0e3, (nt, results["length"], wavelength.size) - ) - results["fit"] = deepcopy(results["spectra"]) - results["texp"] = np.full_like(times, dt) - - results["location"] = np.array([[1.0, 2.0, 3.0]] * results["length"]) - results["direction"] = np.array([[1.0, 2.0, 3.0]] * results["length"]) - - results["element"] = "element" - results["revision"] = np.random.randint(0, 10) - results["x"] = np.random.uniform(Rmin, Rmax, (results["length"],)) - results["y"] = np.random.uniform(Rmin, Rmax, (results["length"],)) - results["z"] = np.random.uniform(zmin, zmax, (results["length"],)) - results["R"] = np.random.uniform(Rmin, Rmax, (results["length"],)) - results["ti"] = np.random.uniform(10, 10.0e3, (nt, results["length"])) - results["vtor"] = np.random.uniform(1.0e2, 1.0e6, (nt, results["length"])) - results["angf"] = np.random.uniform(1.0e2, 1.0e6, (nt, results["length"])) - results["conc"] = np.random.uniform(1.0e-6, 1.0e-1, (nt, results["length"])) - results["bad_channels"] = [] - - for quantity in quantities: - results[f"{quantity}_records"] = [ - f"{quantity}_R_path", - f"{quantity}_z_path", - f"{quantity}_element_path", - f"{quantity}_time_path", - f"{quantity}_ti_path", - f"{quantity}_angf_path", - f"{quantity}_conc_path", - ] - - return results - - def _get_spectrometer( - self, - uid: str, - instrument: str, - revision: RevisionLike, - quantities: Set[str], - ) -> Dict[str, Any]: - - results: Dict[str, Any] = { - "length": np.random.randint(4, 20), - "machine_dims": self.MACHINE_DIMS, - } - dt = np.random.uniform(0.001, 1.0) - times = np.arange(TSTART, TEND, dt) - wavelength = np.arange(520, 530, 0.1) - nt = times.shape[0] - results["times"] = times - results["wavelength"] = wavelength - results["spectra"] = np.random.uniform( - 10, 10.0e3, (nt, results["length"], wavelength.size) - ) - - results["location"] = np.array([[1.0, 2.0, 3.0]] * results["length"]) - results["direction"] = np.array([[1.0, 2.0, 3.0]] * results["length"]) - - results["revision"] = np.random.randint(0, 10) - - for quantity in quantities: - results[f"{quantity}_records"] = [ - f"{quantity}_time_path", - f"{quantity}_spectra_path", - ] - - return results - - def _get_thomson_scattering( - self, - uid: str, - instrument: str, - revision: RevisionLike, - quantities: Set[str], - ) -> Dict[str, Any]: - - Rmin, Rmax = self.MACHINE_DIMS[0][0], self.MACHINE_DIMS[0][1] - zmin, zmax = self.MACHINE_DIMS[1][0], self.MACHINE_DIMS[1][1] - - results: Dict[str, Any] = { - "length": np.random.randint(4, 20), - "machine_dims": self.MACHINE_DIMS, - } - - dt = np.random.uniform(0.001, 1.0) - times = np.arange(TSTART, TEND, dt) - nt = times.shape[0] - results["times"] = times - results["revision"] = np.random.randint(0, 10) - - results["x"] = np.random.uniform(Rmin, Rmax, (results["length"],)) - results["y"] = np.random.uniform(Rmin, Rmax, (results["length"],)) - results["z"] = np.random.uniform(zmin, zmax, (results["length"],)) - results["R"] = np.random.uniform(Rmin, Rmax, (results["length"],)) - results["chi2"] = np.random.uniform(0, 2.0, (nt, results["length"])) - results["te"] = np.random.uniform(10, 10.0e3, (nt, results["length"])) - results["ne"] = np.random.uniform(1.0e16, 1.0e21, (nt, results["length"])) - results["bad_channels"] = [] - - for quantity in quantities: - results[f"{quantity}_records"] = [ - f"{quantity}_Rz_path", - f"{quantity}_value_path", - f"{quantity}_error_path", - ] - - return results - - def _get_equilibrium( - self, - uid: str, - calculation: str, - revision: RevisionLike, - quantities: Set[str], - ) -> Dict[str, Any]: - - Rmin, Rmax = self.MACHINE_DIMS[0][0], self.MACHINE_DIMS[0][1] - zmin, zmax = self.MACHINE_DIMS[1][0], self.MACHINE_DIMS[1][1] - - results: Dict[str, Any] = {} - - dt = np.random.uniform(0.001, 1.0) - times = np.arange(TSTART, TEND, dt) - results["times"] = times - nt = times.shape[0] - nrho = np.random.randint(20, 40) - - results["element"] = "element" - - results["R"] = np.random.uniform(Rmin, Rmax, (nrho,)) - results["z"] = np.random.uniform(zmin, zmax, (nrho,)) - - results["rgeo"] = np.random.uniform(Rmin, Rmax, (nt,)) - results["rmag"] = np.random.uniform(Rmin, Rmax, (nt,)) - results["zmag"] = np.random.uniform(zmin, zmax, (nt,)) - results["ipla"] = np.random.uniform(1.0e4, 1.0e6, (nt,)) - results["wp"] = np.random.uniform(1.0e3, 1.0e5, (nt,)) - results["df"] = np.random.uniform(0, 1, (nt,)) - results["faxs"] = np.random.uniform(1.0e-6, 0.1, (nt,)) - results["fbnd"] = np.random.uniform(-1, 1, (nt,)) - - results["psin"] = np.random.uniform(0, 1, (nrho,)) - results["psi_r"] = np.random.uniform(Rmin, Rmax, (nrho,)) - results["psi_z"] = np.random.uniform(zmin, zmax, (nrho,)) - - results["f"] = np.random.uniform(1.0e-6, 0.1, (nt, nrho)) - results["ftor"] = np.random.uniform(1.0e-4, 1.0e-2, (nt, nrho)) - results["ajac"] = np.random.uniform(1.0e-3, 2.0, (nt, nrho)) - results["vjac"] = np.random.uniform(1.0e-3, 2.0, (nt, nrho)) - results["rmji"] = np.random.uniform(Rmin, Rmax, (nt, nrho)) - results["rmjo"] = np.random.uniform(Rmin, Rmax, (nt, nrho)) - results["rbnd"] = np.random.uniform(Rmin, Rmax, (nt, nrho)) - results["zbnd"] = np.random.uniform(zmin, zmax, (nt, nrho)) - results["psi"] = np.random.uniform(-1, 1, (nt, nrho, nrho)) - - results["revision"] = np.random.randint(0, 10) - - for quantity in quantities: - results[f"{quantity}_records"] = [f"{quantity}_records"] - results["psi_records"] = ["value_records", "R_records", "z_records"] - return results - - def _get_radiation( - self, - uid: str, - instrument: str, - revision: RevisionLike, - quantities: Set[str], - ) -> Dict[str, Any]: - - results: Dict[str, Any] = { - "length": np.random.randint(4, 20), - "machine_dims": self.MACHINE_DIMS, - } - - results["revision"] = np.random.randint(0, 10) - dt = np.random.uniform(0.001, 1.0) - times = np.arange(TSTART, TEND, dt) - results["times"] = times - nt = times.shape[0] - results["location"] = np.array([[1.0, 2.0, 3.0]] * results["length"]) - results["direction"] = np.array([[1.0, 2.0, 3.0]] * results["length"]) - - results["times"] = times - results["brightness"] = np.random.uniform(0, 1.0e6, (nt, results["length"])) - - results["brightness_records"] = ["brightness_records"] * results["length"] - - return results - - def _get_bremsstrahlung_spectroscopy( - self, - uid: str, - instrument: str, - revision: RevisionLike, - quantities: Set[str], - ) -> Dict[str, Any]: - - results: Dict[str, Any] = { - "length": np.random.randint(4, 20), - "machine_dims": self.MACHINE_DIMS, - } - results["revision"] = np.random.randint(0, 10) - dt = np.random.uniform(0.001, 1.0) - times = np.arange(TSTART, TEND, dt) - results["times"] = times - nt = times.shape[0] - - results["location"] = np.array([[1.0, 2.0, 3.0]] * results["length"]) - results["direction"] = np.array([[1.0, 2.0, 3.0]] * results["length"]) - - quantity = "zeff" - results[quantity] = np.random.uniform(0, 1.0e6, (nt, results["length"])) - - results[f"{quantity}_records"] = [ - f"{quantity}_path_records", - f"{quantity}_los_records", - ] - - return results - - def _get_helike_spectroscopy( - self, - uid: str, - instrument: str, - revision: RevisionLike, - quantities: Set[str], - ) -> Dict[str, Any]: - - nwavelength = np.random.randint(256, 1024) - wavelength_start, wavelength_end = 3.8, 4.0 - - results: Dict[str, Any] = { - "length": np.random.randint(4, 20), - "machine_dims": self.MACHINE_DIMS, - } - results["revision"] = np.random.randint(0, 10) - dt = np.random.uniform(0.001, 1.0) - times = np.arange(TSTART, TEND, dt) - results["times"] = times - nt = times.shape[0] - - results["location"] = np.array([[1.0, 2.0, 3.0]] * results["length"]) - results["direction"] = np.array([[1.0, 2.0, 3.0]] * results["length"]) - - results["wavelength"] = np.linspace( - wavelength_start, wavelength_end, nwavelength - ) - - for quantity in quantities: - if quantity == "spectra": - results[quantity] = np.random.uniform( - 0, 1.0e6, (nt, results["length"], nwavelength) - ) - else: - results[quantity] = np.random.uniform(0, 1.0e4, (nt, results["length"])) - - results[f"{quantity}_records"] = [ - f"{quantity}_path_records", - ] - - return results - - def _get_diode_filters( - self, - uid: str, - instrument: str, - revision: RevisionLike, - quantities: Set[str], - ) -> Dict[str, Any]: - - results: Dict[str, Any] = { - "length": np.random.randint(4, 20), - "machine_dims": self.MACHINE_DIMS, - } - - results["revision"] = np.random.randint(0, 10) - dt = np.random.uniform(0.001, 1.0) - times = np.arange(TSTART, TEND, dt) - results["times"] = times - nt = times.shape[0] - - results["location"] = np.array([[1.0, 2.0, 3.0]] * results["length"]) - results["direction"] = np.array([[1.0, 2.0, 3.0]] * results["length"]) - results["labels"] = np.array(["label"] * results["length"]) - - for quantity in quantities: - results[quantity] = np.random.uniform(0, 1.0e6, (nt, results["length"])) - - results[f"{quantity}_records"] = [ - f"{quantity}_path_records", - ] - results[f"{quantity}_error_records"] = [ - f"{quantity}_error_path_records", - ] - - return results - - def _get_interferometry( - self, - uid: str, - instrument: str, - revision: RevisionLike, - quantities: Set[str], - ) -> Dict[str, Any]: - - results: Dict[str, Any] = { - "length": np.random.randint(4, 20), - "machine_dims": self.MACHINE_DIMS, - } - - results["revision"] = np.random.randint(0, 10) - dt = np.random.uniform(0.001, 1.0) - times = np.arange(TSTART, TEND, dt) - results["times"] = times - nt = times.shape[0] - - results["location"] = np.array([[1.0, 2.0, 3.0]] * results["length"]) - results["direction"] = np.array([[1.0, 2.0, 3.0]] * results["length"]) - for quantity in quantities: - results[quantity] = np.random.uniform(0, 1.0e6, (nt, results["length"])) - - results[f"{quantity}_records"] = [ - f"{quantity}_path_records", - ] - results[f"{quantity}_error_records"] = [ - f"{quantity}_error_path_records", - ] - - return results - - def _get_bad_channels( - self, uid: str, instrument: str, quantity: str - ) -> List[Number]: - return [] - - def _set_times_item( - self, - results: Dict[str, Any], - times: np.ndarray, - ): - if "times" not in results: - times = times - - def close(self): - del self._client - - def requires_authentication(self): - return True - - -def _test_get_methods( - instrument="ts", - nsamples=1, -): - """ - Generalised test for all get methods of the abstractreader - """ - - for i in range(nsamples): - reader = Reader( - 1, - TSTART, - TEND, - ) - - quantities = set(AVAILABLE_QUANTITIES[reader.INSTRUMENT_METHODS[instrument]]) - - results = reader.get("", instrument, 0, quantities) - - # Check whether data is as expected - for q, actual, expected in [(q, results[q], results[q]) for q in quantities]: - assert np.all(actual.values == expected) - - -def test_get_thomson_scattering(): - _test_get_methods(instrument="thomson_scattering", nsamples=10) - - -def test_get_charge_exchange(): - _test_get_methods(instrument="charge_exchange", nsamples=10) - - -def test_get_spectrometer(): - _test_get_methods(instrument="spectrometer", nsamples=10) - - -def test_get_equilibrium(): - _test_get_methods(instrument="equilibrium", nsamples=10) - - -def test_get_radiation(): - _test_get_methods(instrument="radiation", nsamples=10) - - -def test_get_bremsstrahlung_spectroscopy(): - _test_get_methods( - instrument="bremsstrahlung_spectroscopy", - nsamples=10, - ) - - -def test_get_helike_spectroscopy(): - _test_get_methods( - instrument="helike_spectroscopy", - nsamples=10, - ) - - -def test_get_diode_filters(): - _test_get_methods( - instrument="diode_filters", - nsamples=10, - ) - - -def test_get_interferometry(): - _test_get_methods( - instrument="interferometry", - nsamples=10, - ) diff --git a/tests/unit/test_offset.py b/tests/unit/test_offset.py deleted file mode 100644 index e69de29b..00000000 diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index 5a076fed..67b3f0ab 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -70,67 +70,6 @@ def test_session_context_manager(): assert session.global_session == new_session assert session.global_session == old_session - -# @given(emails()) -# def test_session_begin(email): -# old_global = session.global_session -# try: -# Session.begin(email) -# assert session.global_session != old_global -# assert session.global_session._user[0].identifier.localpart == email -# finally: -# session.global_session = old_global - - -# @given( -# emails(), lists(from_regex("[a-z0-9]+", fullmatch=True), unique=True, min_size=1) -# ) -# def test_push_pop_agents(email, agent_ids): -# assume(all(ident != email for ident in agent_ids)) -# session = Session(email) -# agents = [] -# original_agent = session._user[0] -# for ident in agent_ids: -# agents.append(session.prov.agent(ident)) -# session.push_agent(agents[-1]) -# assert session.agent is agents[-1] -# for agent in agents[::-1]: -# popped = session.pop_agent() -# assert popped is agent -# delegation = next( -# iter( -# filter( -# lambda x: x.get_attribute("prov:delegate") == {agent.identifier}, -# session.prov.get_records(prov.ProvDelegation), -# ) -# ) -# ) -# assert delegation.get_attribute("prov:responsible") == { -# session.agent.identifier -# } -# assert session.agent is original_agent - - -# @given(emails(), from_regex("[a-z0-9]+", fullmatch=True)) -# def test_new_agent_context(email, agent_id): -# assume(email != agent_id) -# session = Session(email) -# agent = session.prov.agent(agent_id) -# old_agent = session.agent -# with session.new_agent(agent): -# assert session.agent is agent -# delegation = next( -# iter( -# filter( -# lambda x: x.get_attribute("prov:delegate") == {agent.identifier}, -# session.prov.get_records(prov.ProvDelegation), -# ) -# ) -# ) -# assert delegation.get_attribute("prov:responsible") == {old_agent.identifier} -# assert session.agent is old_agent - - def test_dependency_provenance(): session = Session("rand.m.person@ukaea.uk") dependencies = list( @@ -158,8 +97,3 @@ def test_dependency_provenance(): commit = indica_records[0].get_attribute("git_commit") assert len(commit) == 1 assert "UNKNOWN" not in commit - - -# Test exporting/reloading data produces same result - -# Test export/reload properly interacts with data in __main__ From 50278ef8825e7fdc4c100632a5529dadc55986d4 Mon Sep 17 00:00:00 2001 From: Marco Sertoli Date: Fri, 27 Oct 2023 17:50:46 +0100 Subject: [PATCH 07/17] Wrapper to build centrifugal asymmetry parameters and profiles in Indica native DataArray format --- indica/operators/centrifugal_asymmetry.py | 107 ++++++++++++++++++++++ 1 file changed, 107 insertions(+) create mode 100644 indica/operators/centrifugal_asymmetry.py diff --git a/indica/operators/centrifugal_asymmetry.py b/indica/operators/centrifugal_asymmetry.py new file mode 100644 index 00000000..20a31dd8 --- /dev/null +++ b/indica/operators/centrifugal_asymmetry.py @@ -0,0 +1,107 @@ +import matplotlib.pylab as plt +from indica.models.plasma import example_run as example_plasma +from indica.numpy_typing import LabeledArray +from copy import deepcopy + +import numpy as np +from xarray import DataArray + +from indica.datatypes import ELEMENTS +from indica.equilibrium import Equilibrium +import indica.physics as ph +from indica.utilities import assign_data + + +def centrifugal_asymmetry_parameter( + ion_density: DataArray, + ion_temperature: DataArray, + electron_temperature: DataArray, + toroidal_rotation: DataArray, + meanz: DataArray, + zeff: DataArray, + main_ion: str, +): + """ + Indica-native wrapper to calculate the centrifugal asymmetry parameter a-la Wesson + """ + + elements = ion_density.element + zeff = zeff.sum("element") + + asymmetry_parameter = deepcopy(ion_density) + for elem in elements.values: + main_ion_mass = ELEMENTS[main_ion][1] + mass = ELEMENTS[elem][1] + asymmetry_parameter.loc[dict(element=elem)] = ph.centrifugal_asymmetry( + ion_temperature.sel(element=elem).drop_vars("element"), + electron_temperature, + mass, + meanz.sel(element=elem).drop_vars("element"), + zeff, + main_ion_mass, + toroidal_rotation=toroidal_rotation.sel(element=elem).drop_vars("element"), + ) + + return asymmetry_parameter + + +def centrifugal_asymmetry_2d_map( + profile_to_map: DataArray, + asymmetry_parameter: DataArray, + equilibrium: Equilibrium, + t: LabeledArray = None, +): + """Map centrifugal asymmetric profiles to 2D""" + + if t is None: + t = profile_to_map.t.values + + rho_2d = equilibrium.rho.interp(t=t) + R_0 = ( + equilibrium.rmjo.drop_vars("z") + .interp(t=t) + .interp(rho_poloidal=rho_2d) + .drop_vars("rho_poloidal") + ) + + _profile_2d = profile_to_map.interp(rho_poloidal=rho_2d).drop_vars("rho_poloidal") + profile_2d = _profile_2d * np.exp( + asymmetry_parameter.interp(rho_poloidal=rho_2d) * (rho_2d.R ** 2 - R_0 ** 2) + ) + + return profile_2d + + +def example_run(): + + plasma = example_plasma() + + asymmetry_parameter = centrifugal_asymmetry_parameter( + plasma.ion_density, + plasma.ion_temperature, + plasma.electron_temperature, + plasma.toroidal_rotation, + plasma.meanz, + plasma.zeff, + plasma.main_ion, + ) + + ion_density_2d = centrifugal_asymmetry_2d_map( + plasma.ion_density, asymmetry_parameter, plasma.equilibrium, + ) + + ion_density_2d = assign_data( + ion_density_2d, ("density", "ion"), "$m^{-3}$", long_name="Ion density" + ) + + tplot = ion_density_2d.t[2] + element = ion_density_2d.element[2] + plot_2d = ion_density_2d.sel(t=tplot, element=element) + plt.figure() + plot_2d.plot() + + plt.figure() + plot_2d.sel(z=0, method="nearest").plot(label="Midplane z=0") + plt.legend() + + return plasma, ion_density_2d From 9bac848653b0e4985ddc04d28d4fb9c1e3c07163 Mon Sep 17 00:00:00 2001 From: Marco Sertoli Date: Sat, 28 Oct 2023 02:09:30 +0100 Subject: [PATCH 08/17] Slow but working version of inversion including asymmetry --- indica/operators/tomo_asymmetry.py | 183 +++++++++++++++++++++++++++++ 1 file changed, 183 insertions(+) create mode 100644 indica/operators/tomo_asymmetry.py diff --git a/indica/operators/tomo_asymmetry.py b/indica/operators/tomo_asymmetry.py new file mode 100644 index 00000000..d52986c8 --- /dev/null +++ b/indica/operators/tomo_asymmetry.py @@ -0,0 +1,183 @@ +from copy import deepcopy +from typing import Tuple + +import matplotlib.pylab as plt +import numpy as np +from scipy.interpolate import CubicSpline +from scipy.optimize import least_squares +import xarray as xr +from xarray import DataArray + +from indica.converters import LineOfSightTransform +from indica.numpy_typing import ArrayLike +from indica.operators.centrifugal_asymmetry import centrifugal_asymmetry_2d_map + +DataArrayCoords = Tuple[DataArray, DataArray] + + +def InvertPoloidalAsymmetry( + los_integral: DataArray, + los_transform: LineOfSightTransform, + knots: list = [0, 0.2, 0.5, 0.9, 1.0], + xspl: ArrayLike = np.linspace(0, 1.05, 51), + bc_profile: str = "clamped", + bc_asymmetry: str = "natural", + times: ArrayLike = None, +): + """ + Estimates the poloidal distribution from line-of-sight integrals + assuming the local quantity is poloidally asymmetric following + Wesson's formula. + + Parameters + ---------- + los_integral + Measured LOS-integral from e.g. a pinhole camera. + los_transform + Line of sight transform for the forward model (must already have + assigned equilibrium) + knots + The spline knots to use when fitting the emissivity data. + """ + + def residuals(yknots_concat): + yknots_profile = yknots_concat[:nknots] + yknots_asymmetry = yknots_concat[nknots:] + + profile_spline = CubicSpline( + xknots, + yknots_profile, + axis=0, + bc_type=bc_profile, + ) + asymmetry_spline = CubicSpline( + xknots, + yknots_asymmetry, + axis=0, + bc_type=bc_asymmetry, + ) + + profile_to_map = DataArray(profile_spline(xspl), coords=coords) + asymmetry_parameter = DataArray(asymmetry_spline(xspl), coords=coords) + + profile_2d = centrifugal_asymmetry_2d_map( + profile_to_map, + asymmetry_parameter, + equilibrium=los_transform.equilibrium, + t=t, + ) + _bckc = los_transform.integrate_on_los(profile_2d, t=t) + return (_data - _bckc) / _error + + if times is None: + times = los_integral.t + + if hasattr(los_integral, "error"): + error = los_integral.error + else: + error = los_integral * 0.1 + + coords = [("rho_poloidal", xspl)] + + nknots = np.size(knots) + xknots = np.array(knots) + # xknots_concat = np.append(np.array(knots), np.array(knots)) + + guess = los_integral.isel(t=0).mean().values + yknots_emission = np.full_like(xknots, guess) + yknots_asymmetry = np.full_like(xknots, 0) + yknots_concat = np.append(yknots_emission, yknots_asymmetry) + + lower_bound, upper_bound = set_bounds(xknots) + + profile = xr.DataArray( + np.empty((np.size(xspl), np.size(times))), + coords=[("rho_poloidal", xspl), ("t", times)], + ) + asymmetry = deepcopy(profile) + profile_2d = [] + bckc = xr.full_like(los_integral, np.nan) + for t in times: + _data = los_integral.sel(t=t) + _error = error.sel(t=t) + + fit = least_squares( + residuals, + yknots_concat, + bounds=(lower_bound, upper_bound), + verbose=True, + ) + + profile_spline = CubicSpline( + xknots, + fit.x[:nknots], + axis=0, + bc_type=bc_profile, + ) + asymmetry_spline = CubicSpline( + xknots, + fit.x[nknots:], + axis=0, + bc_type=bc_asymmetry, + ) + + profile.loc[dict(t=t)] = profile_spline(xspl) + asymmetry.loc[dict(t=t)] = asymmetry_spline(xspl) + + _profile_2d = centrifugal_asymmetry_2d_map( + profile, + asymmetry, + equilibrium=los_transform.equilibrium, + t=t, + ) + bckc.loc[dict(t=t)] = los_transform.integrate_on_los(_profile_2d, t=t) + profile_2d.append(_profile_2d) + + profile_2d = xr.concat(profile_2d, "t") + + return profile_2d, bckc, profile, asymmetry + + +def set_bounds(xknots): + l_bound_profile = np.full_like(xknots, 0) + u_bound_profile = np.full_like(xknots, 1) + l_bound_asymmetry = np.full_like(xknots, -10) + u_bound_asymmetry = np.full_like(xknots, 10) + l_bound_asymmetry[0] = 0 + u_bound_asymmetry[0] = 1 + + lower_bound = np.append(l_bound_profile, l_bound_asymmetry) + upper_bound = np.append(u_bound_profile, u_bound_asymmetry) + + return lower_bound, upper_bound + + +def example_run(): + import indica.models.bolometer_camera as bolo + import indica.operators.centrifugal_asymmetry as centrifugal_asymmetry + + plasma, ion_density_2d = centrifugal_asymmetry.example_run() + _, model, _ = bolo.example_run() + + profile_2d = ion_density_2d.sel(element="ar") + los_transform = model.los_transform + los_integral = los_transform.integrate_on_los(profile_2d, profile_2d.t.values) + + times = los_integral.t.values[:4] + norm_factor = los_integral.max().values + profile_2d_bckc, bckc, profile_bckc, asymmetry_bckc = InvertPoloidalAsymmetry( + los_integral / norm_factor, los_transform, times=times + ) + profile_2d_bckc, bckc = profile_2d_bckc * norm_factor, bckc * norm_factor + + for t in times: + plt.ioff() + plt.figure() + los_integral.sel(t=t).plot(marker="o") + bckc.sel(t=t).plot() + plt.figure() + ((profile_2d.sel(t=t) - profile_2d_bckc.sel(t=t)) / profile_2d.sel(t=t)).plot() + plt.figure() + profile_2d.sel(t=t).sel(z=0, method="nearest").plot(marker="o") + profile_2d_bckc.sel(t=t).sel(z=0, method="nearest").plot() + plt.show() From 4e46bb96f33bc87fc46d39a5e507b5bc31121b9b Mon Sep 17 00:00:00 2001 From: Marco Sertoli Date: Sat, 28 Oct 2023 02:10:31 +0100 Subject: [PATCH 09/17] Making tangential as default bolometer geometry. TODO: make general geometry sets to be used across models. --- indica/models/bolometer_camera.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/indica/models/bolometer_camera.py b/indica/models/bolometer_camera.py index 6dff3798..1a680af4 100644 --- a/indica/models/bolometer_camera.py +++ b/indica/models/bolometer_camera.py @@ -174,13 +174,21 @@ def example_run( # Create new interferometers diagnostics if origin is None or direction is None: los_end = np.full((nchannels, 3), 0.0) - los_end[:, 0] = 0.17 - los_end[:, 1] = 0.0 - los_end[:, 2] = np.linspace(0.6, -0.6, nchannels) - los_start = np.array([[1.0, 0, 0]] * los_end.shape[0]) + los_end[:, 0] = 0.0 + los_end[:, 1] = np.linspace(-0.2, -1, nchannels) + los_end[:, 2] = 0.0 + los_start = np.array([[1.5, 0, 0]] * los_end.shape[0]) origin = los_start direction = los_end - los_start + # los_end = np.full((nchannels, 3), 0.0) + # los_end[:, 0] = 0.17 + # los_end[:, 1] = 0.0 + # los_end[:, 2] = np.linspace(0.6, -0.6, nchannels) + # los_start = np.array([[1.0, 0, 0]] * los_end.shape[0]) + # origin = los_start + # direction = los_end - los_start + los_transform = LineOfSightTransform( origin[:, 0], origin[:, 1], From ec394968f4a88ee76dd246806881e6a9095a720d Mon Sep 17 00:00:00 2001 From: Marco Sertoli Date: Sat, 28 Oct 2023 02:11:22 +0100 Subject: [PATCH 10/17] Modified centrifugal_asymmetry.py to accept profiles without t dimension --- indica/operators/centrifugal_asymmetry.py | 25 ++++++++++++++++------- 1 file changed, 18 insertions(+), 7 deletions(-) diff --git a/indica/operators/centrifugal_asymmetry.py b/indica/operators/centrifugal_asymmetry.py index 20a31dd8..51a45f13 100644 --- a/indica/operators/centrifugal_asymmetry.py +++ b/indica/operators/centrifugal_asymmetry.py @@ -1,13 +1,13 @@ -import matplotlib.pylab as plt -from indica.models.plasma import example_run as example_plasma -from indica.numpy_typing import LabeledArray from copy import deepcopy +import matplotlib.pylab as plt import numpy as np from xarray import DataArray from indica.datatypes import ELEMENTS from indica.equilibrium import Equilibrium +from indica.models.plasma import example_run as example_plasma +from indica.numpy_typing import LabeledArray import indica.physics as ph from indica.utilities import assign_data @@ -56,7 +56,16 @@ def centrifugal_asymmetry_2d_map( if t is None: t = profile_to_map.t.values - rho_2d = equilibrium.rho.interp(t=t) + if "t" in profile_to_map.dims: + _profile_to_map = profile_to_map.sel(t=t, method="nearest") + else: + _profile_to_map = profile_to_map + if "t" in asymmetry_parameter.dims: + _asymmetry_parameter = asymmetry_parameter.sel(t=t, method="nearest") + else: + _asymmetry_parameter = asymmetry_parameter + + rho_2d = equilibrium.rho.sel(t=t, method="nearest") R_0 = ( equilibrium.rmjo.drop_vars("z") .interp(t=t) @@ -64,9 +73,9 @@ def centrifugal_asymmetry_2d_map( .drop_vars("rho_poloidal") ) - _profile_2d = profile_to_map.interp(rho_poloidal=rho_2d).drop_vars("rho_poloidal") + _profile_2d = _profile_to_map.interp(rho_poloidal=rho_2d).drop_vars("rho_poloidal") profile_2d = _profile_2d * np.exp( - asymmetry_parameter.interp(rho_poloidal=rho_2d) * (rho_2d.R ** 2 - R_0 ** 2) + _asymmetry_parameter.interp(rho_poloidal=rho_2d) * (rho_2d.R**2 - R_0**2) ) return profile_2d @@ -87,7 +96,9 @@ def example_run(): ) ion_density_2d = centrifugal_asymmetry_2d_map( - plasma.ion_density, asymmetry_parameter, plasma.equilibrium, + plasma.ion_density, + asymmetry_parameter, + plasma.equilibrium, ) ion_density_2d = assign_data( From 6c2cc0725dfabd807d40c13f6adcbde3ff82af3c Mon Sep 17 00:00:00 2001 From: Marco Sertoli Date: Sat, 28 Oct 2023 02:11:58 +0100 Subject: [PATCH 11/17] Residuals must be residuals, with a sign --- indica/operators/spline_fit_easy.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/indica/operators/spline_fit_easy.py b/indica/operators/spline_fit_easy.py index d514abd3..a9e1520b 100644 --- a/indica/operators/spline_fit_easy.py +++ b/indica/operators/spline_fit_easy.py @@ -29,8 +29,7 @@ def residuals(yknots): bc_type=bc_type, ) bckc = np.interp(x, xspl, spline(xspl)) - residuals = np.sqrt((y - bckc) ** 2 / err**2) - return residuals + return (y - bckc) / err yspl = xr.DataArray( np.empty((len(xspl), len(ydata.t))), From 16f2c15f477a6891876b954c50ebc1a945ccab10 Mon Sep 17 00:00:00 2001 From: Marco Sertoli Date: Thu, 2 Nov 2023 11:10:04 +0000 Subject: [PATCH 12/17] Removed tomo_asymmetry.py to be pushed to a separate branch --- indica/operators/tomo_asymmetry.py | 183 ----------------------------- 1 file changed, 183 deletions(-) delete mode 100644 indica/operators/tomo_asymmetry.py diff --git a/indica/operators/tomo_asymmetry.py b/indica/operators/tomo_asymmetry.py deleted file mode 100644 index d52986c8..00000000 --- a/indica/operators/tomo_asymmetry.py +++ /dev/null @@ -1,183 +0,0 @@ -from copy import deepcopy -from typing import Tuple - -import matplotlib.pylab as plt -import numpy as np -from scipy.interpolate import CubicSpline -from scipy.optimize import least_squares -import xarray as xr -from xarray import DataArray - -from indica.converters import LineOfSightTransform -from indica.numpy_typing import ArrayLike -from indica.operators.centrifugal_asymmetry import centrifugal_asymmetry_2d_map - -DataArrayCoords = Tuple[DataArray, DataArray] - - -def InvertPoloidalAsymmetry( - los_integral: DataArray, - los_transform: LineOfSightTransform, - knots: list = [0, 0.2, 0.5, 0.9, 1.0], - xspl: ArrayLike = np.linspace(0, 1.05, 51), - bc_profile: str = "clamped", - bc_asymmetry: str = "natural", - times: ArrayLike = None, -): - """ - Estimates the poloidal distribution from line-of-sight integrals - assuming the local quantity is poloidally asymmetric following - Wesson's formula. - - Parameters - ---------- - los_integral - Measured LOS-integral from e.g. a pinhole camera. - los_transform - Line of sight transform for the forward model (must already have - assigned equilibrium) - knots - The spline knots to use when fitting the emissivity data. - """ - - def residuals(yknots_concat): - yknots_profile = yknots_concat[:nknots] - yknots_asymmetry = yknots_concat[nknots:] - - profile_spline = CubicSpline( - xknots, - yknots_profile, - axis=0, - bc_type=bc_profile, - ) - asymmetry_spline = CubicSpline( - xknots, - yknots_asymmetry, - axis=0, - bc_type=bc_asymmetry, - ) - - profile_to_map = DataArray(profile_spline(xspl), coords=coords) - asymmetry_parameter = DataArray(asymmetry_spline(xspl), coords=coords) - - profile_2d = centrifugal_asymmetry_2d_map( - profile_to_map, - asymmetry_parameter, - equilibrium=los_transform.equilibrium, - t=t, - ) - _bckc = los_transform.integrate_on_los(profile_2d, t=t) - return (_data - _bckc) / _error - - if times is None: - times = los_integral.t - - if hasattr(los_integral, "error"): - error = los_integral.error - else: - error = los_integral * 0.1 - - coords = [("rho_poloidal", xspl)] - - nknots = np.size(knots) - xknots = np.array(knots) - # xknots_concat = np.append(np.array(knots), np.array(knots)) - - guess = los_integral.isel(t=0).mean().values - yknots_emission = np.full_like(xknots, guess) - yknots_asymmetry = np.full_like(xknots, 0) - yknots_concat = np.append(yknots_emission, yknots_asymmetry) - - lower_bound, upper_bound = set_bounds(xknots) - - profile = xr.DataArray( - np.empty((np.size(xspl), np.size(times))), - coords=[("rho_poloidal", xspl), ("t", times)], - ) - asymmetry = deepcopy(profile) - profile_2d = [] - bckc = xr.full_like(los_integral, np.nan) - for t in times: - _data = los_integral.sel(t=t) - _error = error.sel(t=t) - - fit = least_squares( - residuals, - yknots_concat, - bounds=(lower_bound, upper_bound), - verbose=True, - ) - - profile_spline = CubicSpline( - xknots, - fit.x[:nknots], - axis=0, - bc_type=bc_profile, - ) - asymmetry_spline = CubicSpline( - xknots, - fit.x[nknots:], - axis=0, - bc_type=bc_asymmetry, - ) - - profile.loc[dict(t=t)] = profile_spline(xspl) - asymmetry.loc[dict(t=t)] = asymmetry_spline(xspl) - - _profile_2d = centrifugal_asymmetry_2d_map( - profile, - asymmetry, - equilibrium=los_transform.equilibrium, - t=t, - ) - bckc.loc[dict(t=t)] = los_transform.integrate_on_los(_profile_2d, t=t) - profile_2d.append(_profile_2d) - - profile_2d = xr.concat(profile_2d, "t") - - return profile_2d, bckc, profile, asymmetry - - -def set_bounds(xknots): - l_bound_profile = np.full_like(xknots, 0) - u_bound_profile = np.full_like(xknots, 1) - l_bound_asymmetry = np.full_like(xknots, -10) - u_bound_asymmetry = np.full_like(xknots, 10) - l_bound_asymmetry[0] = 0 - u_bound_asymmetry[0] = 1 - - lower_bound = np.append(l_bound_profile, l_bound_asymmetry) - upper_bound = np.append(u_bound_profile, u_bound_asymmetry) - - return lower_bound, upper_bound - - -def example_run(): - import indica.models.bolometer_camera as bolo - import indica.operators.centrifugal_asymmetry as centrifugal_asymmetry - - plasma, ion_density_2d = centrifugal_asymmetry.example_run() - _, model, _ = bolo.example_run() - - profile_2d = ion_density_2d.sel(element="ar") - los_transform = model.los_transform - los_integral = los_transform.integrate_on_los(profile_2d, profile_2d.t.values) - - times = los_integral.t.values[:4] - norm_factor = los_integral.max().values - profile_2d_bckc, bckc, profile_bckc, asymmetry_bckc = InvertPoloidalAsymmetry( - los_integral / norm_factor, los_transform, times=times - ) - profile_2d_bckc, bckc = profile_2d_bckc * norm_factor, bckc * norm_factor - - for t in times: - plt.ioff() - plt.figure() - los_integral.sel(t=t).plot(marker="o") - bckc.sel(t=t).plot() - plt.figure() - ((profile_2d.sel(t=t) - profile_2d_bckc.sel(t=t)) / profile_2d.sel(t=t)).plot() - plt.figure() - profile_2d.sel(t=t).sel(z=0, method="nearest").plot(marker="o") - profile_2d_bckc.sel(t=t).sel(z=0, method="nearest").plot() - plt.show() From ef590a61c626bc8bac96d762e47a6738870aef0f Mon Sep 17 00:00:00 2001 From: Marco Sertoli Date: Thu, 2 Nov 2023 11:39:48 +0000 Subject: [PATCH 13/17] Moved surf_los.dat in data folder and ran pre-commit --- indica/{readers => data}/surf_los.dat | 0 indica/readers/ppfreader.py | 2 +- tests/unit/readers/test_surf.py | 2 +- tests/unit/test_session.py | 1 + 4 files changed, 3 insertions(+), 2 deletions(-) rename indica/{readers => data}/surf_los.dat (100%) diff --git a/indica/readers/surf_los.dat b/indica/data/surf_los.dat similarity index 100% rename from indica/readers/surf_los.dat rename to indica/data/surf_los.dat diff --git a/indica/readers/ppfreader.py b/indica/readers/ppfreader.py index 0f2a44a1..ccf563bf 100644 --- a/indica/readers/ppfreader.py +++ b/indica/readers/ppfreader.py @@ -33,7 +33,7 @@ from ..utilities import to_filename -SURF_PATH = Path(surf_los.__file__).parent / "surf_los.dat" +SURF_PATH = Path(surf_los.__file__).parent.parent / "data/surf_los.dat" class PPFError(Exception): diff --git a/tests/unit/readers/test_surf.py b/tests/unit/readers/test_surf.py index e98810f4..0a7c6ba5 100644 --- a/tests/unit/readers/test_surf.py +++ b/tests/unit/readers/test_surf.py @@ -20,7 +20,7 @@ import indica.readers.surf_los as surf_los -filepath = Path(surf_los.__file__).parent / "surf_los.dat" +filepath = Path(surf_los.__file__).parent.parent / "data/surf_los.dat" PIXEL = 0.00099 diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index 67b3f0ab..ed5f3a09 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -70,6 +70,7 @@ def test_session_context_manager(): assert session.global_session == new_session assert session.global_session == old_session + def test_dependency_provenance(): session = Session("rand.m.person@ukaea.uk") dependencies = list( From e1a8c3965380827a67f8239da7ef502bc867ec1e Mon Sep 17 00:00:00 2001 From: Marco Sertoli Date: Thu, 2 Nov 2023 11:45:56 +0000 Subject: [PATCH 14/17] Changed name of spline_fit_easy.py to simply spline_fit.py --- indica/operators/{spline_fit_easy.py => spline_fit.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename indica/operators/{spline_fit_easy.py => spline_fit.py} (100%) diff --git a/indica/operators/spline_fit_easy.py b/indica/operators/spline_fit.py similarity index 100% rename from indica/operators/spline_fit_easy.py rename to indica/operators/spline_fit.py From b2deee581356a4a04711c8570a7a76a1d92b6e6a Mon Sep 17 00:00:00 2001 From: Marco Sertoli Date: Thu, 2 Nov 2023 11:54:33 +0000 Subject: [PATCH 15/17] Reinstated asymmetry tomography because of dependencies. --- indica/operators/tomo_asymmetry.py | 224 +++++++++++++++++++++++++++++ 1 file changed, 224 insertions(+) create mode 100644 indica/operators/tomo_asymmetry.py diff --git a/indica/operators/tomo_asymmetry.py b/indica/operators/tomo_asymmetry.py new file mode 100644 index 00000000..eb57d410 --- /dev/null +++ b/indica/operators/tomo_asymmetry.py @@ -0,0 +1,224 @@ +from copy import deepcopy +from typing import Tuple + +import matplotlib.pylab as plt +import numpy as np +from scipy.interpolate import CubicSpline +from scipy.optimize import least_squares +import xarray as xr +from xarray import DataArray + +from indica.converters import LineOfSightTransform +from indica.numpy_typing import ArrayLike, LabeledArray +from indica.operators.centrifugal_asymmetry import centrifugal_asymmetry_2d_map + +DataArrayCoords = Tuple[DataArray, DataArray] + + +class InvertPoloidalAsymmetry: + def __init__( + self, + knots: LabeledArray = [0, 0.2, 0.5, 0.9, 1.0], + xspl: ArrayLike = np.linspace(0, 1.05, 51), + bounds_profile: tuple = (0, np.inf), + bounds_asymmetry: tuple = (-np.inf, np.inf), + bc_profile: str = "clamped", + bc_asymmetry: str = "natural", + ): + + # xknots must include 0, 1, and rho_max + rho_max = np.max(xspl) + _xknots = np.append(np.array(knots), [0.0, 1.0, rho_max]) + _xknots = np.unique(np.sort(_xknots)) + xknots = _xknots[np.where((_xknots >= -0) * (_xknots <= rho_max))[0]] + + # Profile knots to scan <= 1. + mask_profile = (xknots >= 0) * (xknots <= 1.0) + indx_profile = np.where(mask_profile)[0] + nknots_profile = len(indx_profile) + + # Asymmetry knots to scan >0 and <1. + mask_asymmetry = (xknots > 0) * (xknots < 1.0) + indx_asymmetry = np.where(mask_asymmetry)[0] + indx_profile[-1] + nknots_asymmetry = len(indx_asymmetry) + + # Define optimisation bounds + lbounds = np.append( + [bounds_profile[0]] * nknots_profile, + [bounds_asymmetry[0]] * nknots_asymmetry, + ) + ubounds = np.append( + [bounds_profile[1]] * nknots_profile, + [bounds_asymmetry[1]] * nknots_asymmetry, + ) + bounds = (lbounds, ubounds) + + self.xspl = xspl + self.xknots = xknots + self.nknots = np.size(xknots) + self.bounds = bounds + self.mask_profile = mask_profile + self.mask_asymmetry = mask_asymmetry + self.indx_profile = indx_profile + self.indx_asymmetry = indx_asymmetry + self.bc_profile = bc_profile + self.bc_asymmetry = bc_asymmetry + + def __call__( + self, + los_integral: DataArray, + los_transform: LineOfSightTransform, + t: LabeledArray = None, + debug: bool = False, + ): + """ + Estimates the poloidal distribution from line-of-sight integrals + assuming the local quantity is poloidally asymmetric following + Wesson's formula. + + Parameters + ---------- + los_integral + Measured LOS-integral from e.g. a pinhole camera. + los_transform + Line of sight transform for the forward model (must already have + assigned equilibrium) + knots + The spline knots to use when fitting the emissivity data. + """ + + coords = [("rho_poloidal", self.xspl)] + mask_profile, mask_asymmetry = self.mask_profile, self.mask_asymmetry + indx_profile, indx_asymmetry = self.indx_profile, self.indx_asymmetry + xknots, xspl = self.xknots, self.xspl + equilibrium = los_transform.equilibrium + bc_profile, bc_asymmetry = self.bc_profile, self.bc_asymmetry + + def evaluate(xknots, yconcat): + """ + 1. Separate profile and asymmetry knots + add boundaries + 2. Initialize and call splines. + 3. Map profile and asymmetry to 2D + """ + yprofile = yconcat[indx_profile] + yprofile = np.append(yprofile, 0.0) + yasymmetry = yconcat[indx_asymmetry] + yasymmetry = np.append(yasymmetry[0], yasymmetry) + yasymmetry = np.append(yasymmetry, [yasymmetry[-1], yasymmetry[-1]]) + + profile_spline = CubicSpline(xknots, yprofile, 0, bc_profile,) + profile_to_map = DataArray(profile_spline(xspl), coords=coords) + asymmetry_spline = CubicSpline(xknots, yasymmetry, 0, bc_asymmetry,) + asymmetry_parameter = DataArray(asymmetry_spline(xspl), coords=coords) + profile_2d = centrifugal_asymmetry_2d_map( + profile_to_map, asymmetry_parameter, equilibrium=equilibrium, t=t, + ) + return profile_2d, profile_to_map, asymmetry_parameter + + def residuals(yknots_concat): + """ + Calculate the residuals to minimize + """ + profile_2d, _, _ = evaluate(xknots, yknots_concat) + _bckc = los_transform.integrate_on_los(profile_2d, t=t) + return (_data - _bckc) / _error + + if t is None: + t = los_integral.t + else: + t = los_integral.t.sel(t=t, method="nearest") + + if hasattr(los_integral, "error"): + error = los_integral.error + else: + error = los_integral * 0.1 + + self.t = np.array(t, ndmin=1) + self.los_transform = los_transform + norm_factor = los_integral.max().values + self.data = los_integral / norm_factor + self.error = error / norm_factor + + # Initial guesses + _guess_profile = self.data.sel(t=self.t[0]).mean().values + _guess_asymmetry = 0.0 + guess_profile = np.full_like(xknots, _guess_profile) + guess_asymmetry = np.full_like(xknots, _guess_asymmetry) + yconcat = np.append(guess_profile[mask_profile], guess_asymmetry[mask_asymmetry],) + + profile = [] + asymmetry = [] + profile_2d = [] + bckc = [] + for t in self.t: + _data = self.data.sel(t=t).values + _error = self.error.sel(t=t).values + _x = np.arange(len(_data)) + + if debug: + plt.figure() + plt.ioff() + plt.errorbar(_x, _data, _error, marker="o") + + fit = least_squares(residuals, yconcat, bounds=self.bounds, verbose=debug,) + yconcat = fit.x + + _profile_2d, _profile_to_map, _asymmetry_parameter = evaluate(xknots, yconcat) + _bckc = los_transform.integrate_on_los(_profile_2d, t=t) + if debug: + plt.plot(_x, _bckc, linewidth=2) + plt.show() + + profile.append(_profile_to_map) + asymmetry.append(_asymmetry_parameter) + profile_2d.append(_profile_2d) + bckc.append(_bckc) + + profile = xr.concat(profile, "t") * norm_factor + asymmetry = xr.concat(asymmetry, "t") + profile_2d = xr.concat(profile_2d, "t") * norm_factor + bckc = xr.concat(bckc, "t") * norm_factor + + # xr.DataArray( + # np.empty((np.size(self.xspl), np.size(self.t))), + # coords=[("rho_poloidal", self.xspl), ("t", self.t)], + # ) + + return profile_2d, bckc, profile, asymmetry + + +def example_run(): + import indica.models.bolometer_camera as bolo + import indica.operators.centrifugal_asymmetry as centrifugal_asymmetry + from indica.operators.tomo_asymmetry_class import InvertPoloidalAsymmetry + + plasma, ion_density_2d = centrifugal_asymmetry.example_run() + _, model, _ = bolo.example_run() + + profile_2d = ion_density_2d.sel(element="ar") + los_transform = model.los_transform + los_integral = los_transform.integrate_on_los(profile_2d, profile_2d.t.values) + + print("\n Running asymmetry inference...") + self = InvertPoloidalAsymmetry() + + profile_2d_bckc, bckc, profile, asymmetry = self( + los_integral, los_transform, + ) + print("...and finished \n") + + for t in bckc.t: + plt.ioff() + plt.figure() + los_integral.sel(t=t).plot(marker="o") + bckc.sel(t=t).plot() + plt.figure() + ((profile_2d.sel(t=t) - profile_2d_bckc.sel(t=t)) / profile_2d.sel(t=t)).plot() + plt.figure() + profile_2d.sel(t=t).sel(z=0, method="nearest").plot(marker="o") + profile_2d_bckc.sel(t=t).sel(z=0, method="nearest").plot() + plt.show() + + +if __name__ == "__main__": + example_run() From 69c9c5ec8359fbc38db4dd48a5777a3aec16ccaf Mon Sep 17 00:00:00 2001 From: Marco Sertoli Date: Thu, 2 Nov 2023 12:31:13 +0000 Subject: [PATCH 16/17] Reinstated asymmetry tomography because of dependencies and cleaned up implementation. --- indica/operators/centrifugal_asymmetry.py | 19 ++-- indica/operators/tomo_asymmetry.py | 100 ++++++++++++++++------ 2 files changed, 86 insertions(+), 33 deletions(-) diff --git a/indica/operators/centrifugal_asymmetry.py b/indica/operators/centrifugal_asymmetry.py index 51a45f13..f929c9e7 100644 --- a/indica/operators/centrifugal_asymmetry.py +++ b/indica/operators/centrifugal_asymmetry.py @@ -81,7 +81,7 @@ def centrifugal_asymmetry_2d_map( return profile_2d -def example_run(): +def example_run(plot: bool = False): plasma = example_plasma() @@ -105,14 +105,15 @@ def example_run(): ion_density_2d, ("density", "ion"), "$m^{-3}$", long_name="Ion density" ) - tplot = ion_density_2d.t[2] - element = ion_density_2d.element[2] - plot_2d = ion_density_2d.sel(t=tplot, element=element) - plt.figure() - plot_2d.plot() + if plot: + tplot = ion_density_2d.t[2] + element = ion_density_2d.element[2] + plot_2d = ion_density_2d.sel(t=tplot, element=element) + plt.figure() + plot_2d.plot() - plt.figure() - plot_2d.sel(z=0, method="nearest").plot(label="Midplane z=0") - plt.legend() + plt.figure() + plot_2d.sel(z=0, method="nearest").plot(label="Midplane z=0") + plt.legend() return plasma, ion_density_2d diff --git a/indica/operators/tomo_asymmetry.py b/indica/operators/tomo_asymmetry.py index eb57d410..bd3fcb8f 100644 --- a/indica/operators/tomo_asymmetry.py +++ b/indica/operators/tomo_asymmetry.py @@ -1,4 +1,3 @@ -from copy import deepcopy from typing import Tuple import matplotlib.pylab as plt @@ -9,7 +8,10 @@ from xarray import DataArray from indica.converters import LineOfSightTransform -from indica.numpy_typing import ArrayLike, LabeledArray +import indica.models.bolometer_camera as bolo +from indica.numpy_typing import ArrayLike +from indica.numpy_typing import LabeledArray +import indica.operators.centrifugal_asymmetry as centrifugal_asymmetry from indica.operators.centrifugal_asymmetry import centrifugal_asymmetry_2d_map DataArrayCoords = Tuple[DataArray, DataArray] @@ -97,6 +99,7 @@ def __call__( def evaluate(xknots, yconcat): """ 1. Separate profile and asymmetry knots + add boundaries + asymmetry(rho>=1) 2. Initialize and call splines. 3. Map profile and asymmetry to 2D """ @@ -106,12 +109,25 @@ def evaluate(xknots, yconcat): yasymmetry = np.append(yasymmetry[0], yasymmetry) yasymmetry = np.append(yasymmetry, [yasymmetry[-1], yasymmetry[-1]]) - profile_spline = CubicSpline(xknots, yprofile, 0, bc_profile,) + profile_spline = CubicSpline( + xknots, + yprofile, + 0, + bc_profile, + ) profile_to_map = DataArray(profile_spline(xspl), coords=coords) - asymmetry_spline = CubicSpline(xknots, yasymmetry, 0, bc_asymmetry,) + asymmetry_spline = CubicSpline( + xknots, + yasymmetry, + 0, + bc_asymmetry, + ) asymmetry_parameter = DataArray(asymmetry_spline(xspl), coords=coords) profile_2d = centrifugal_asymmetry_2d_map( - profile_to_map, asymmetry_parameter, equilibrium=equilibrium, t=t, + profile_to_map, + asymmetry_parameter, + equilibrium=equilibrium, + t=t, ) return profile_2d, profile_to_map, asymmetry_parameter @@ -135,24 +151,31 @@ def residuals(yknots_concat): self.t = np.array(t, ndmin=1) self.los_transform = los_transform + self.data = los_integral + self.error = error + + # Normalise to have similar ranges for profile and asymmetry norm_factor = los_integral.max().values - self.data = los_integral / norm_factor - self.error = error / norm_factor + data = los_integral / norm_factor + error = error / norm_factor # Initial guesses - _guess_profile = self.data.sel(t=self.t[0]).mean().values + _guess_profile = data.sel(t=self.t[0]).mean().values _guess_asymmetry = 0.0 guess_profile = np.full_like(xknots, _guess_profile) guess_asymmetry = np.full_like(xknots, _guess_asymmetry) - yconcat = np.append(guess_profile[mask_profile], guess_asymmetry[mask_asymmetry],) + yconcat = np.append( + guess_profile[mask_profile], + guess_asymmetry[mask_asymmetry], + ) profile = [] asymmetry = [] profile_2d = [] bckc = [] for t in self.t: - _data = self.data.sel(t=t).values - _error = self.error.sel(t=t).values + _data = data.sel(t=t).values + _error = error.sel(t=t).values _x = np.arange(len(_data)) if debug: @@ -160,10 +183,17 @@ def residuals(yknots_concat): plt.ioff() plt.errorbar(_x, _data, _error, marker="o") - fit = least_squares(residuals, yconcat, bounds=self.bounds, verbose=debug,) + fit = least_squares( + residuals, + yconcat, + bounds=self.bounds, + verbose=debug, + ) yconcat = fit.x - _profile_2d, _profile_to_map, _asymmetry_parameter = evaluate(xknots, yconcat) + _profile_2d, _profile_to_map, _asymmetry_parameter = evaluate( + xknots, yconcat + ) _bckc = los_transform.integrate_on_los(_profile_2d, t=t) if debug: plt.plot(_x, _bckc, linewidth=2) @@ -188,9 +218,6 @@ def residuals(yknots_concat): def example_run(): - import indica.models.bolometer_camera as bolo - import indica.operators.centrifugal_asymmetry as centrifugal_asymmetry - from indica.operators.tomo_asymmetry_class import InvertPoloidalAsymmetry plasma, ion_density_2d = centrifugal_asymmetry.example_run() _, model, _ = bolo.example_run() @@ -200,23 +227,48 @@ def example_run(): los_integral = los_transform.integrate_on_los(profile_2d, profile_2d.t.values) print("\n Running asymmetry inference...") - self = InvertPoloidalAsymmetry() + invert_asymm = InvertPoloidalAsymmetry() - profile_2d_bckc, bckc, profile, asymmetry = self( - los_integral, los_transform, + profile_2d_bckc, bckc, profile, asymmetry = invert_asymm( + los_integral, + los_transform, ) print("...and finished \n") for t in bckc.t: plt.ioff() + + kwargs_phantom = dict( + marker="o", + label="Phantom", + color="k", + zorder=1, + ) + kwargs_bckc = dict( + marker="x", label="Bckc", color="b", zorder=2, linewidth=4, alpha=0.5 + ) plt.figure() - los_integral.sel(t=t).plot(marker="o") - bckc.sel(t=t).plot() + plt.errorbar( + invert_asymm.data.channel, + invert_asymm.data.sel(t=t), + invert_asymm.error.sel(t=t), + **kwargs_phantom, + ) + bckc.sel(t=t).plot(**kwargs_bckc) + plt.title("LOS integrals") + plt.figure() - ((profile_2d.sel(t=t) - profile_2d_bckc.sel(t=t)) / profile_2d.sel(t=t)).plot() + profile_2d.sel(t=t).plot() + plt.title("2D phantom") + + plt.figure() + profile_2d_bckc.sel(t=t).plot() + plt.title("2D bckc") + plt.figure() - profile_2d.sel(t=t).sel(z=0, method="nearest").plot(marker="o") - profile_2d_bckc.sel(t=t).sel(z=0, method="nearest").plot() + profile_2d.sel(t=t).sel(z=0, method="nearest").plot(**kwargs_phantom) + profile_2d_bckc.sel(t=t).sel(z=0, method="nearest").plot(**kwargs_bckc) + plt.title("Midplane cut (z=0)") plt.show() From bc9b84232035b42a36fd4b6a3033c01989a31324 Mon Sep 17 00:00:00 2001 From: Marco Sertoli Date: Fri, 3 Nov 2023 10:53:31 +0000 Subject: [PATCH 17/17] Minor correction to plotting. --- indica/operators/tomo_asymmetry.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/indica/operators/tomo_asymmetry.py b/indica/operators/tomo_asymmetry.py index bd3fcb8f..2a3c7eac 100644 --- a/indica/operators/tomo_asymmetry.py +++ b/indica/operators/tomo_asymmetry.py @@ -235,6 +235,8 @@ def example_run(): ) print("...and finished \n") + los_transform.plot() + for t in bckc.t: plt.ioff() @@ -256,6 +258,7 @@ def example_run(): ) bckc.sel(t=t).plot(**kwargs_bckc) plt.title("LOS integrals") + plt.legend() plt.figure() profile_2d.sel(t=t).plot() @@ -269,6 +272,7 @@ def example_run(): profile_2d.sel(t=t).sel(z=0, method="nearest").plot(**kwargs_phantom) profile_2d_bckc.sel(t=t).sel(z=0, method="nearest").plot(**kwargs_bckc) plt.title("Midplane cut (z=0)") + plt.legend() plt.show()